loss_func = torch.nn.MSELoss()
loss_his =[[],[],[],[]]# 记录损失for epoch inrange(EPOCH):for step,(batch_x, batch_y)inenumerate(loader):for net, opt, l_his inzip(nets, optimizers, loss_his):
output = net(batch_x)# get output for every net
loss = loss_func(output, batch_y)# compute loss for every net
opt.zero_grad()# clear gradients for next train
loss.backward()# backpropagation, compute gradients
opt.step()# apply gradients
l_his.append(loss.data.numpy())# loss recoder
labels =['SGD','Momentum','RMSprop','Adam']
6.可视化结果
for i, l_his inenumerate(loss_his):
plt.plot(l_his, label=labels[i])
plt.legend(loc='best')
plt.xlabel('Steps')
plt.ylabel('Loss')
plt.ylim((0,0.2))
plt.show()