pytorch生成一个简单的线性模型

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）

1.准备数据集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]])

2.使用类设计模型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

3.构造损失函数和优化器Construct loss and optimizer

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)

4.训练循环Training cycle

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())

运行结果