PyTorch 框架实现线性回归：从数据预处理到模型训练全流程

系列文章目录

01-PyTorch新手必看：张量是什么？5 分钟教你快速创建张量！
02-张量运算真简单！PyTorch 数值计算操作完全指南
03-Numpy 还是 PyTorch？张量与 Numpy 的神奇转换技巧
04-揭秘数据处理神器：PyTorch 张量拼接与拆分实用技巧
05-深度学习从索引开始：PyTorch 张量索引与切片最全解析
06-张量形状任意改！PyTorch reshape、transpose 操作超详细教程
07-深入解读 PyTorch 张量运算：6 大核心函数全面解析，代码示例一步到位！
08-自动微分到底有多强？PyTorch 自动求导机制深度解析
09-从零手写线性回归模型：PyTorch 实现深度学习入门教程
10-PyTorch 框架实现线性回归：从数据预处理到模型训练全流程

文章目录

系列文章目录
前言
一、构建数据集
- 1.1 示例代码
二、数据加载器
- 2.1 示例代码
三、定义模型
- 3.1 示例代码
四、定义损失函数与优化器
- 4.1 示例代码
五、训练模型
- 5.1 示例代码
六、绘制结果
- 6.1 示例代码
七、主函数
- 7.1 示例代码
- 7.2 示例输出
八、总结
- 8.1 完整代码

前言

在之前的文章中，通过手动方式构建了一个简单的线性回归模型。然而，面对复杂的网络设计，手动实现不仅繁琐，还容易出错。为了提高效率和灵活性，可以利用 PyTorch 提供的组件来快速搭建模型。

本文将通过 PyTorch 实现线性回归，主要包括以下内容：

使用 nn.MSELoss 代替自定义的平方损失函数
使用 data.DataLoader 代替自定义的数据加载器
使用 optim.SGD 代替自定义的优化器
使用 nn.Linear 代替自定义的线性模型

一、构建数据集

首先，需要生成一组模拟的线性回归数据，并将其转换为 PyTorch 张量。使用 sklearn.datasets.make_regression 函数来创建数据。

1.1 示例代码

import torch
from sklearn.datasets import make_regression

# 构建数据集
def create_dataset():
    x, y, coef = make_regression(
        n_samples=150,       # 样本数量
        n_features=1,        # 特征数量
        noise=15,            # 噪声大小
        coef=True,           # 返回系数
        bias=10.0,           # 偏置项
        random_state=42      # 随机种子
    )
    # 转换为 PyTorch 张量
    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)
    return x, y, coef

在上面的代码中，生成了一个 150 个样本、带有噪声的数据集，并将其转换为 PyTorch 支持的张量格式，便于后续训练使用。

二、数据加载器

为了更方便地处理数据批量，使用 PyTorch 的 DataLoader 来加载数据集。

2.1 示例代码

from torch.utils.data import TensorDataset, DataLoader

# 构建数据加载器
def create_dataloader(x, y):
    dataset = TensorDataset(x, y)  # 创建数据集对象
    dataloader = DataLoader(dataset, batch_size=20, shuffle=True)  # 批量大小为20
    return dataloader

DataLoader 可以轻松实现数据的分批次处理，同时支持打乱数据顺序以提高模型的泛化能力。

三、定义模型

PyTorch 提供了 nn.Linear 作为线性模型的实现，只需定义输入和输出的特征数量。

线性模型的核心公式为：
$\hat{y} = x \cdot w + b$
其中，w 是权重，b 是偏置，均为模型的可学习参数。

3.1 示例代码

import torch.nn as nn

# 构建线性模型
def create_model():
    model = nn.Linear(in_features=1, out_features=1)  # 输入和输出特征均为1
    return model

四、定义损失函数与优化器

使用 nn.MSELoss 作为损失函数，并使用 optim.SGD 优化器来更新模型参数。

4.1 示例代码

import torch.optim as optim

# 定义损失函数和优化器
def create_loss_and_optimizer(model):
    criterion = nn.MSELoss()  # 均方误差损失
    optimizer = optim.SGD(model.parameters(), lr=0.01)  # 学习率为0.01
    return criterion, optimizer

五、训练模型

将数据加载器、模型、损失函数和优化器整合到训练循环中，通过梯度下降来优化模型。

5.1 示例代码

# 训练模型
def train_model(dataloader, model, criterion, optimizer, epochs=100):
    for epoch in range(epochs):
        for batch_x, batch_y in dataloader:
            # 前向传播：计算预测值
            y_pred = model(batch_x)
            # 计算损失
            loss = criterion(y_pred, batch_y.unsqueeze(1))
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播：计算梯度
            loss.backward()
            # 更新参数
            optimizer.step()
        # 打印每10个epoch的损失
        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}")

在训练过程中，每次从数据加载器中取出一个批次的数据进行训练，并更新模型参数。

六、绘制结果

训练完成后，可视化模型的拟合效果，与真实数据进行对比。

6.1 示例代码

# 可视化模型拟合结果
def plot_results(x, y, model, coef, bias):
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
    plt.rcParams['axes.unicode_minus'] = False   # 正常显示负号

    # 绘制散点图
    plt.scatter(x.numpy(), y.numpy(), label="数据点", alpha=0.7)

    # 绘制模型预测的直线
    x_range = torch.linspace(x.min(), x.max(), 100).reshape(-1, 1)
    y_pred = model(x_range).detach().numpy()
    plt.plot(x_range.numpy(), y_pred, label="拟合直线", color="r")

    # 绘制真实的直线
    coef = torch.tensor(coef, dtype=torch.float32)
    bias = torch.tensor(bias, dtype=torch.float32)
    y_true = coef * x_range + bias
    plt.plot(x_range.numpy(), y_true.numpy(), label="真实直线", color="g", linestyle="--")

    # 添加标题和标签
    plt.title("线性回归拟合结果")
    plt.xlabel("自变量 X")
    plt.ylabel("因变量 Y")

    # 显示图例和网格
    plt.legend()
    plt.grid(True)

    # 显示绘图
    plt.show()

七、主函数

7.1 示例代码

if __name__ == "__main__":
    # 构建数据
    x, y, coef = create_dataset()
    dataloader = create_dataloader(x, y)
    model = create_model()
    criterion, optimizer = create_loss_and_optimizer(model)
    
    # 训练模型
    train_model(dataloader, model, criterion, optimizer, epochs=100)
    
    # 绘制结果
    plot_results(x, y, model, coef, bias=10.0)

7.2 示例输出

在这里插入图片描述
运行结果将显示数据点与拟合直线，直线与真实线性关系高度吻合，说明模型训练效果良好。

八、总结

本文通过使用 PyTorch 实现线性回归模型，完成了以下内容：

构建数据集并设计数据加载器；
使用 nn.Linear 定义线性假设函数；
使用 nn.MSELoss 设计均方误差损失函数；
使用 optim.SGD 实现随机梯度下降优化方法；
训练模型并可视化拟合结果。

8.1 完整代码

import torch
from torch.utils.data import TensorDataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_regression
import matplotlib.pyplot as plt


# 构建数据集
def create_dataset():
    x, y, coef = make_regression(
        n_samples=150,       # 样本数量
        n_features=1,        # 特征数量
        noise=15,            # 噪声大小
        coef=True,           # 返回系数
        bias=10.0,           # 偏置项
        random_state=42      # 随机种子
    )
    # 转换为 PyTorch 张量
    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)
    return x, y, coef

# 构建数据加载器
def create_dataloader(x, y):
    dataset = TensorDataset(x, y)  # 创建数据集对象
    dataloader = DataLoader(dataset, batch_size=20, shuffle=True)  # 批量大小为20
    return dataloader

# 构建线性模型
def create_model():
    model = nn.Linear(in_features=1, out_features=1)  # 输入和输出特征均为1
    return model

# 定义损失函数和优化器
def create_loss_and_optimizer(model):
    criterion = nn.MSELoss()  # 均方误差损失
    optimizer = optim.SGD(model.parameters(), lr=0.01)  # 学习率为0.01
    return criterion, optimizer

# 训练模型
def train_model(dataloader, model, criterion, optimizer, epochs=100):
    for epoch in range(epochs):
        for batch_x, batch_y in dataloader:
            # 前向传播：计算预测值
            y_pred = model(batch_x)
            # 计算损失
            loss = criterion(y_pred, batch_y.unsqueeze(1))
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播：计算梯度
            loss.backward()
            # 更新参数
            optimizer.step()
        # 打印每10个epoch的损失
        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}")

# 可视化模型拟合结果
def plot_results(x, y, model, coef, bias):
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
    plt.rcParams['axes.unicode_minus'] = False   # 正常显示负号

    # 绘制散点图
    plt.scatter(x.numpy(), y.numpy(), label="数据点", alpha=0.7)

    # 绘制模型预测的直线
    x_range = torch.linspace(x.min(), x.max(), 100).reshape(-1, 1)
    y_pred = model(x_range).detach().numpy()
    plt.plot(x_range.numpy(), y_pred, label="拟合直线", color="r")

    # 绘制真实的直线
    coef = torch.tensor(coef, dtype=torch.float32)
    bias = torch.tensor(bias, dtype=torch.float32)
    y_true = coef * x_range + bias
    plt.plot(x_range.numpy(), y_true.numpy(), label="真实直线", color="g", linestyle="--")

    # 添加标题和标签
    plt.title("线性回归拟合结果")
    plt.xlabel("自变量 X")
    plt.ylabel("因变量 Y")

    # 显示图例和网格
    plt.legend()
    plt.grid(True)

    # 显示绘图
    plt.show()



# 主函数
if __name__ == "__main__":
    # 构建数据
    x, y, coef = create_dataset()
    dataloader = create_dataloader(x, y)
    model = create_model()
    criterion, optimizer = create_loss_and_optimizer(model)
    
    # 训练模型
    train_model(dataloader, model, criterion, optimizer, epochs=100)
    
    # 绘制结果
    plot_results(x, y, model, coef, bias=10.0)