PyTorch 的线性回归代码模板

预测结果类似如下，直线是模型

然后代码模板如下，直接复制修改即可，注释也写的很详细了

# 导入 PyTorch 库
import torch 
import torch.nn as nn 
# 导入 NumPy 库，用于数据处理
import numpy as np
# 导入 matplotlib 库，用于数据可视化
import matplotlib.pyplot as plt
# 导入 PyTorch 中的 Variable 类，用于自动梯度计算
from torch.autograd import Variable

# 设置模型输入和输出的维度
input_size = 1
output_size = 1
# 设置学习率
learning_rate = 0.001

# 设置训练数据
xtrain = np.array([...], dtype=np.float32)
ytrain = np.array([...], dtype=np.float32)

# 可视化训练数据
plt.figure()
plt.scatter(xtrain,ytrain) 
plt.xlabel('xtrain') #设置 x 轴标签
plt.ylabel('ytrain') #设置 y 轴标签
plt.show() #显示图像

# 定义线性回归模型
class LinearRegression(nn.Module):
    def __init__(self, input_size, output_size): 
        super(LinearRegression, self).__init__()
        # 定义线性层
        self.linear = nn.Linear(input_size, output_size)  
    def forward(self, x): 
        # 前向传播
        out = self.linear(x) 
        return out

# 定义损失函数，使用均方误差损失
criterion = nn.MSELoss()

# 初始化模型
model = LinearRegression(input_size,output_size)

# 定义优化器，使用随机梯度下降算法
optimizer = torch.optim.SGD(model.parameters(),lr=learning_rate) 

# 设置训练轮数
num_epochs=1000

# 训练模型
for epoch in range(num_epochs):
    # 将 NumPy 数组转化为 PyTorch 的 Variable
    inputs = Variable(torch.from_numpy(xtrain))
    targets = Variable(torch.from_numpy(ytrain))

    # 清零梯度
    optimizer.zero_grad()
    
    # 前向传播
    outputs = model(inputs)
    
    # 计算损失
    loss = criterion(outputs, targets)
    
    # 反向传播
    loss.backward()
    
    # 更新参数
    optimizer.step()
    
    # 每 50 轮打印一次损失值
    if (epoch+1) % 50 == 0:
        print('Epoch [%d/%d], Loss: %.4f' %(epoch+1, num_epochs, loss.item())) 

# 评估模式
model.eval()

# 使用模型进行预测
predicted = model(Variable(torch.from_numpy(xtrain))).data.numpy()

# 可视化预测结果
plt.plot(xtrain,ytrain,'ro')
plt.plot(xtrain,predicted,label='predict')
plt.legend()
plt.show()