用PyTorch实现一个简单的线性回归模型

一、实现思路

线性回归的核心是拟合 $y = wx + b$ 这个线性关系，我们会通过以下步骤实现：

1. 生成模拟的线性数据（含少量噪声，贴近真实场景）；
2. 定义线性回归模型（仅包含权重 $w$ 和偏置 $b$）；
3. 选择损失函数（均方误差MSE）和优化器（梯度下降Adam）；
4. 训练模型，迭代更新 $w$ 和 $b$；
5. 可视化训练结果，验证拟合效果。

二、完整代码实现

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

# 设置随机种子，保证结果可复现
torch.manual_seed(42)
np.random.seed(42)

# ===================== 1. 生成模拟数据 =====================
# 生成x数据：-10到10之间的100个均匀分布的数，形状为(100, 1)
x = torch.linspace(-10, 10, 100).reshape(-1, 1)
# 真实的线性关系：y = 2x + 5，添加少量正态分布噪声
y_true = 2 * x + 5 + torch.normal(0, 1, size=x.shape)

# ===================== 2. 定义线性回归模型 =====================
class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        # 定义线性层：输入维度1，输出维度1（对应y = wx + b）
        self.linear = nn.Linear(in_features=1, out_features=1)
    
    # 前向传播（必须实现的方法）
    def forward(self, x):
        return self.linear(x)

# 实例化模型
model = LinearRegression()

# ===================== 3. 定义损失函数和优化器 =====================
# 损失函数：均方误差（MSE），适合回归任务
criterion = nn.MSELoss()
# 优化器：Adam（自适应学习率，比普通梯度下降更稳定），学习率0.01
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# ===================== 4. 训练模型 =====================
epochs = 200  # 训练轮数
loss_history = []  # 记录每轮损失，方便后续可视化

for epoch in range(epochs):
    # 训练模式（线性回归无dropout/bn，可省略，但养成习惯）
    model.train()
    
    # 1. 前向传播：用模型预测y
    y_pred = model(x)
    
    # 2. 计算损失：预测值与真实值的MSE
    loss = criterion(y_pred, y_true)
    
    # 3. 反向传播：计算梯度
    optimizer.zero_grad()  # 梯度清零（避免累加）
    loss.backward()        # 反向传播求梯度
    
    # 4. 更新参数：用优化器更新w和b
    optimizer.step()
    
    # 记录损失
    loss_history.append(loss.item())
    
    # 每20轮打印一次训练信息
    if (epoch + 1) % 20 == 0:
        # 提取当前的w和b（model.linear.weight是张量，需转成数值）
        w = model.linear.weight.item()
        b = model.linear.bias.item()
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}, w: {w:.4f}, b: {b:.4f}")

# ===================== 5. 模型评估与可视化 =====================
# 切换到评估模式
model.eval()
# 禁用梯度计算（节省内存，加速推理）
with torch.no_grad():
    y_pred_final = model(x)

# 打印最终拟合的参数
final_w = model.linear.weight.item()
final_b = model.linear.bias.item()
print("\n最终拟合的线性方程：y = {:.4f}x + {:.4f}".format(final_w, final_b))

# 可视化：原始数据 + 拟合直线
plt.figure(figsize=(10, 6))
# 绘制原始数据点
plt.scatter(x.numpy(), y_true.numpy(), label="Raw Data", color="blue", alpha=0.6)
# 绘制拟合直线
plt.plot(x.numpy(), y_pred_final.numpy(), label="Fitted Line", color="red", linewidth=2)
plt.xlabel("x")
plt.ylabel("y")
plt.title("Linear Regression with PyTorch")
plt.legend()
plt.grid(True)
plt.show()

# 可视化损失变化
plt.figure(figsize=(10, 6))
plt.plot(loss_history, label="Training Loss", color="green")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.title("Loss Trend During Training")
plt.legend()
plt.grid(True)
plt.show()

三、代码关键部分解释

1. 数据生成：

   • 用torch.linspace生成均匀分布的x，reshape(-1, 1)将一维数据转为二维（适配PyTorch线性层的输入要求）；
   • 真实标签y_true基于$y=2x+5$生成，添加torch.normal的噪声，模拟真实场景的非完美线性数据。

2. 模型定义：

   • 继承nn.Module（PyTorch所有模型的基类）；
   • nn.Linear(1, 1)是核心，内部封装了权重$w$和偏置$b$，无需手动定义参数，简化代码。

3. 训练流程：

   • optimizer.zero_grad()：每次迭代前清空梯度，否则梯度会累加导致参数更新错误；
   • loss.backward()：自动计算所有可训练参数的梯度；
   • optimizer.step()：用计算好的梯度更新$w$和$b$。

4. 评估阶段：

   • model.eval()：切换到评估模式（对线性回归无实质影响，但处理含Dropout/BatchNorm的模型时必须）；
   • with torch.no_grad()：禁用梯度计算，减少内存占用，加快推理速度。