Python线性回归模型实现——up主：刘二大人《PyTorch深度学习实践》

b站up主：刘二大人《PyTorch深度学习实践》
教程： https://www.bilibili.com/video/BV1Y7411d7Ys?p=5&vd_source=715b347a0d6cb8aa3822e5a102f366fe
$$\begin{gathered} 线性回归模型：:torch.nn.Linear \\ 损失函数：nn.MSELoss \\ 优化器：optim.Adam \end{gathered}$$

import torch
import matplotlib.pyplot as plt

#数据集
x_data = torch.Tensor([[1.0],[2.0],[3.0]])
y_data = torch.Tensor([[2.0],[4.0],[6.0]])

epoch_p = []
loss_p = []

#设计模型
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

model = LinearModel()

#构造损失函数和优化器
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.01)

#train
for epoch in range(1000):
  y_pred = model(x_data)
  loss = criterion(y_pred, y_data)
  loss_p.append(loss.item())
  epoch_p.append(epoch)
  print("第{}个Epoch，loss = {:.6f}".format(epoch, loss))

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

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

plt.figure()
plt.plot(epoch_p, loss_p, c = 'b')
plt.xlabel('Epoch')
plt.ylabel('loss')
plt.show()