使用pytorch实现一个简单的线性模型demo

使用pytorch实现一个简单的线性模型demo

模型从建立到训练的一个完整的步骤

1. prepare dataset
2. design model using Class（inherit from nn.Moudle）
3. Construct loss and optimizer（using Pytorch API）
4. Training cycle（forward backward and update）

下面以一个简单的线性模型为例：

prepare dataset

# 导入pytorch包
import torch

# 简单的数据集和对应的label
x_data = torch.tensor([[1.0], [2.0], [3.0]])
y_data = torch.tensor([[2.0], [4.0], [6.0]])

design model using Class

class LinearModel(torch.nn.Module):
    # 构造函数
    def __init__(self):
        super(LinearModel, self).__init__()
        # Linear （in_features: int, out_features: int, bias: bool = True, device=None, dtype=None)
        # 实例化对象
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        # linear是个对象，对象后面加括号，表示创建了一个可调用的对象（callerable）
        y_pred = self.linear(x)
        return y_pred

Construct loss and optimizer（using Pytorch API）

model = LinearModel()
# criterion = torch.nn.MSELoss(size_average=False)
# 使用MSE损失（均方误差损失）作为loss损失
criterion = torch.nn.MSELoss(reduction='sum')
# 优化器使用SGD优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

Training cycle（forward backward and update）

for epoch in range(1000):
    y_pred = model(x_data)
    # loss是个对象,criterion也是一个对象，对象后面直接接参数相当于调用当前对象的__call__方法
    # 对象后面直接跟参数的前提是这个实例化对象对应的类实现了__call__方法
    loss = criterion(y_pred, y_data)
    print(epoch, loss.item())
    # 梯度置0
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 自动更新，根据梯度，所设置的学习率自动更新
    optimizer.step()

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

完整代码

import torch


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


class LinearModel(torch.nn.Module):
    # 构造函数
    def __init__(self):
        super(LinearModel, self).__init__()
        # Linear （in_features: int, out_features: int, bias: bool = True, device=None, dtype=None)
        # 实例化对象
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        # linear是个对象，对象后面加括号，表示创建了一个可调用的对象（callable）
        y_pred = self.linear(x)
        return y_pred


model = LinearModel()
# criterion = torch.nn.MSELoss(size_average=False)
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(1000):
    y_pred = model(x_data)
    # loss是个对象
    loss = criterion(y_pred, y_data)
    # 打印loss时会自动调用loss.__str()__
    print(epoch, loss.item())
    # 梯度置0
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 自动更新，根据梯度，所设置的学习率自动更新
    optimizer.step()

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

最终结果