pytorch回归网络

import torch
import torch.optim
import matplotlib.pyplot as plt
import torch.nn.functional as f
from torch.autograd import Variable

# unsequeeze(param) 增加param的维度，在torch中处理2维的张量
x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
# 目标函数表达式
y = x.pow(2)+0.2*torch.rand(x.size())
# torch中计算张量时需将数据放置Variable中
x, y = Variable(x), Variable(y)

# plt.scatter(x.data.numpy(), y.data.numpy())
# plt.show()

# 建立一个（feature=1，hidden=10，output=1）的网络，并且在hidden使用激励函数
net2 = torch.nn.Sequential(
    torch.nn.Linear(1, 10),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 1)
)

# 选择优化器SGD:随机梯度下降。优化的参数时网络中的参数，lr学习效率是0.2
optimizer = torch.optim.SGD(net2.parameters(), lr=0.2)
# 选择计算损失函数，在回归问题中一般使用MSELoss()
loss_func = torch.nn.MSELoss()

# 训练100步
for i in range(100):
    predict = net2(x)
    # 计算损失值
    loss = loss_func(predict, y)

    # 清除之前优化器中的权重
    optimizer.zero_grad()
    # 损失值反向传播给网络
    loss.backward()
    # 按步优化
    optimizer.step()
    # 每五步画图一次
    if i % 5 == 0:
        # plot and show learning process
        plt.cla()
        plt.scatter(x.data.numpy(), y.data.numpy())
        plt.plot(x.data.numpy(), predict.data.numpy(), 'r-', lw=5)
        plt.text(0.5, 0, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'red'})
        plt.pause(0.1)
plt.ioff()
plt.show()