使用神经网络拟合argmax函数

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_19332527/article/details/84075212

最近比较闲，听说神经网络可以拟合任意函数，因此尝试使用神经网络拟合argmax函数。

我使用的是pytorch，就是堆叠多层感知机来拟合一个argmax函数。为了防止不收敛，我的输入是一个 20维的向量，要求输出的值是 0~1 范围之内的，0表示第0个数是最大的，1表示第19个数是最大的，0.4表示第15个数是最大的，可以用  int(输出*20) 转回下标。代码如下：

import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.optim import SGD

#定义多层感知机
class Net(nn.Module):
	def __init__(self):
		super(Net,self).__init__()
		self.l1=nn.Linear(20,512)
		self.relu=nn.ReLU()
		self.l2=nn.Linear(512,1)
		self.sig=nn.Sigmoid()
	def forward(self,x):
		x=self.relu(self.l1(x))
		x=self.sig(self.l2(x))
		return x

#随机生成训练集
def get_data():
	x=Variable(torch.randn(4000,20))
	_,y=torch.max(x,dim=1)
	y=y.float()/20
	return x,y
#训练函数定义
def train(step=1000,lr=1e-3):
	model=Net()
	criterion=nn.MSELoss() #损失函数使用均方误差函数
	optimizer=SGD(model.parameters(),lr=lr)
	step=step
	x,y=get_data() #获取训练集
	for i in range(step):
		y_=model(x)
		loss=criterion(y_,y)
		
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

		print("{} step, loss is {}".format(i,loss.data[0]))
	torch.save(model.state_dict(),"./model.pth") #最后保存模型
#测试模型
def test(x,weights=None):
	model=Net()
	#加载参数
    if weights is not None:
		model.load_state_dict(torch.load(weights))
	y_=model(x)  #进行预测
	_,y=torch.max(x,dim=1) #求出最大值下标
	y=y.float()/20  #最大值下标除以20，才能和y_进行比较
	print(y,y_) #输出y和y_进行大概的比较

if __name__=="__main__":
	train(10000,0.1)
	x=Variable((torch.randn(5,20))) #随机生成测试集
	test(x,"./model.pth")

最后我是训练了10000轮，得出的结果如下所示，好像不是很好，估计是训练还不够充分或者模型比较简单，可以加大数据量或者增加模型复杂度:

最后结果