pytorch笔记3-简单线性回归

使用pytorch构造简单线性模型

import torch
import matplotlib.pyplot as plt
x_data = torch.tensor([[1.0], [2.0], [3.0]])#构建3*1的矩阵
y_data = torch.tensor([[2.0], [4.0], [6.0]])

class LinearModel(torch.nn.Module):  #创建类构造模型，最少包含下面两个函数
    def __init__(self):
        super(LinearModel, self).__init__()
        # 该线性层需要学习的参数是w和b  获取w、b的方式是linear.weight、linear.bias
        self.linear = torch.nn.Linear(1, 1) #构造Linear对象，包含w和b，自动完成w*x+b的计算和反向传播，(1,1)是x和y的特征维度

    def forward(self, x): #正向传播，callable
        y_pred = self.linear(x)
        return y_pred

model = LinearModel() #callable，可y=model（x）

criterion = torch.nn.MSELoss(size_average=True) #创建损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 参数优化，创建对象，model.parameters（）自动传入所有参数
lepoch=[]
lloss=[]

for epoch in range(100): #循环训练，不断优化参数
    y_pred = model(x_data)  # 正向计算预测值
    loss = criterion(y_pred, y_data)  # 计算loss
    print(epoch, loss.item())
    lepoch.append(epoch)
    lloss.append(loss.item())
    optimizer.zero_grad()  # 梯度清0
    loss.backward()  # 反向传播，自动计算梯度
    optimizer.step()  # 更新参数，即更新w和b的值

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

x_test = torch.tensor([[4.0]])
y_test = model(x_test)
print('y_pred = ', y_test.data)
plt.plot(lepoch,lloss)
plt.show()