Pytorch实现线性回归模型

文章目录

• ​​任务​​
• ​​实现​​

任务

根据给出的训练集，训练一个模型，拟合一条直线

预期效果：

实现

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import numpy as np
from torch import

# 训练数据集
x_train = np.array([[3.3],[4.4],[5.5],[6.6],[6.69],[4.4],[9.8],[6.12],[7.7],[2.67],[7.42],[10.91],[5.13],[7.97],[3.1]], dtype=np.float32)
y_train = np.array([[1.7],[2.76],[2.09],[3.19],[1.694],[1.573],[3.366],[2.596],[2.53],[1.221],[2.827],[3.465],[1.65],[2.904],[1.3]], dtype=np.float32)

plt.plot(x_train, y_train, 'ro')
plt.show()

# 将训练数据集转为tensor
x_train = torch.from_numpy(x_train)
y_train = torch.from_numpy(y_train)

# 定义一个线性模型（神经网络）
model = nn.Linear(1,1)
# 优化器 完成模型参数和权重更新
optim = optim.SGD(model.parameters(),lr=0.01)
# 损失函数
Loss = MSELoss()

#训练过程
ecophs = 30000
for i in range(ecophs):
  # 计算模型输出
  outputs = model.forward(x_train)
  # 计算损失
  loss = Loss(y_train,outputs)
  # 在梯度反向传播之前，梯度降为0
  model.zero_grad()
  # 反向传播
  loss.backward()
  # 优化器优化
  optim.step()

# 获得训练的参数
w,a = model.parameters()
print(w.item(),a.item())

# 可视化
predicted = model.forward(x_train).data.numpy()
plt.plot(x_train,y_train,'ro')
plt.plot(x_train,predicted)