基于pytorch搭建一元线性回归模型

'''
基于pytorch搭建一元线性回归模型
'''
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt


class LinerRegression(nn.Module):
    def __init__(self):
        super(LinerRegression, self).__init__()
        self.linear = nn.Linear(1, 1)  # 输入和输出均为一维

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


x_train = np.array(
    [[3.3], [4.4], [5.5], [6.71], [6.93], [4.168], [9.779], [6.182], [7.59], [2.167], [7.042], [10.791], [5.313],
     [7.997], [3.1]], dtype=np.float32)
y_train = np.array(
    [[1.7], [2.76], [2.09], [3.19], [1.694], [1.573], [3.366], [2.596], [2.53], [1.221], [2.827], [3.465], [1.65],
     [2.904], [1.3]], dtype=np.float32)


x_train = torch.from_numpy(x_train)
y_train = torch.from_numpy(y_train)

model = LinerRegression()
criterion = nn.MSELoss()  # 定义损失函数和优化器，均方差作为优化器

num_epochs = 1000
for epoch in range(num_epochs):
    input = x_train
    target = y_train
    out = model(input)  # 训练，得到输出
    loss = criterion(out, target)  # 计算损失

    optimizer = optim.SGD(model.parameters(), lr=1e-3)  # 随机梯度下降优化器
    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()  # 没有预测值，这里用eval把值固定
predict = model(x_train)  # 得到训练结果
predict = predict.data.numpy()

plt.plot(x_train.numpy(), y_train.numpy(), 'ro', label='Original_data')
plt.plot(x_train.numpy(), predict, label='Fitting Line')
plt.show()