pytorch实现简单线性模型

一、数据集的预处理

x_data = torch.tensor([[1.0], [2.0], [3.0]])
y_data = torch.tensor([[2.0], [4.0], [6.0]])

使用 torch.tensor 函数对数据进行预处理。

二、线性模型的构建

class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

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

1. 用类的形式定义线性模型，继承自 torch.nn.Module。
2. 构建函数使用 super 表示调用父类的构建函数。
3. 线性模型使用 torch.nn.Linear，参数（1，1）表示输入维数为1， 输出维数为1。
4. forword 函数为前向传播函数，命名一定为 forward， 是对父类同名函数的覆盖。

三、参数设定

1. 实例化线性模型。

model = LinearModel()

2. 构建损失函数。MSE为平均平方误差。size_average=False 表示最后的损失不用处以平均数M。

criterion = torch.nn.MSELoss(size_average=False)

3. 优化器的设置。使用SGD的方式来求最小值，model.parameters() 表示将线性模型所需要优化参数传入，lr为学习率。

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

四、迭代求最优参数

for epoch in range(1000):
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    print(epoch, loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print('w =', model.linear.weight.item())
print('b =', model.linear.bias.item())

• 设置迭代的次数为1000
• 用线性model求预测值
• 求损失的值
• 每次迭代将参数的梯度值归零
• 求反向传播
• 更新参数

五、打印结果

x_test = torch.Tensor([4.0])
y_test = model(x_test)
print('y_pred', y_test.data.item())