pytorch 线性模型

from matplotlib import pyplot as plt
import torch
from torch import nn

X = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
Y = 4*X + 5 + torch.rand(X.size())


class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.liner = nn.Linear(1, 1)    # 输入输出的维度都是 1

    def forward(self, x):
        out = self.liner(x)
        return out


model = LinearRegression()
model.cuda()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)

num_epochs = 1000
for epoch in range(num_epochs):
    inputs = torch.autograd.Variable(X).cuda()
    target = torch.autograd.Variable(Y).cuda()

    # 向前传播
    out = model(inputs)
    loss = criterion(out, target)

    # 向后传播
    optimizer.zero_grad()   # 注意每次迭代都需要清零
    loss.backward()
    optimizer.step()    # 更新所有参数

    if (epoch + 1) % 20 == 0:
        print('Epoch[{}/{}], loss:{:.6f}'.format(epoch + 1, num_epochs, loss.item()))

model.eval().cuda()
predict = model(torch.autograd.Variable(X).cuda())
predict = predict.data.cpu().numpy()
plt.plot(X.numpy(), Y.numpy(), 'ro', label='Original Data')
plt.plot(X.numpy(), predict, label='Fitting Line')
plt.show()