预测结果类似如下,直线是模型
然后代码模板如下,直接复制修改即可,注释也写的很详细了
# 导入 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()