PyTorch:nn 模块

3 阅读3分钟

PyTorch:nn 模块

本示例实现一个三阶多项式模型,通过最小化欧式距离平方和,拟合区间 ([-π, π]) 内的正弦函数 (y=\sin(x))。

该实现的核心变化:

  • 使用 PyTorch 的 nn 包构建神经网络(高层 API)
  • nn 包封装了各类神经网络层(Module),无需手动定义参数和前向传播
  • 仍基于 autograd 完成自动微分,但无需手动管理参数和梯度

nn 包的核心价值:

  • autograd 适合定义计算图和求梯度,但偏底层
  • nn 包提供模块化的神经网络层,每个 Module 包含可训练参数,能接收输入并生成输出
  • 极大简化复杂神经网络的构建过程

运行输出

99 1022.1802978515625
199 686.7949829101562
299 462.7146301269531
399 312.9090881347656
499 212.69467163085938
599 145.6101531982422
699 100.6719741821289
799 70.54719543457031
899 50.337425231933594
999 36.76863479614258
1099 27.651264190673828
1199 21.519664764404297
1299 17.39257049560547
1399 14.612120628356934
1499 12.73715591430664
1599 11.471559524536133
1699 10.616497039794922
1799 10.038167953491211
1899 9.646638870239258
1999 9.381280899047852
Result: y = -0.016341309994459152 + 0.8391954898834229 x + 0.0028191469609737396 x^2 + -0.09083471447229385 x^3

完整代码(带详细注释)

import torch
import math

# 创建输入和输出张量
x = torch.linspace(-math.pi, math.pi, 2000)  # [-π, π] 生成2000个均匀分布的点
y = torch.sin(x)  # 正弦函数真实值

# 构造特征张量:将 x 转换为 (x, x², x³) 特征
# 因为三阶多项式 y = a + bx + cx² + dx³ 是 (x, x², x³) 的线性组合
p = torch.tensor([1, 2, 3])  # 幂次:1, 2, 3
xx = x.unsqueeze(-1).pow(p)  # 维度变换+广播:(2000,) → (2000,1) → (2000,3)

# 解释维度变换:
# x.unsqueeze(-1) → shape (2000, 1)
# p → shape (3,)
# 广播机制使两者相乘得到 shape (2000, 3) 的特征矩阵

# 使用 nn 包构建模型:Sequential 按顺序堆叠神经网络层
# Linear(3, 1):线性层,输入维度3(x, x², x³),输出维度1(预测值)
# Flatten(0, 1):展平层,将 (2000,1) 展平为 (2000,),匹配 y 的形状
model = torch.nn.Sequential(
    torch.nn.Linear(3, 1),    # 线性层:y = w1*x + w2*x² + w3*x³ + b
    torch.nn.Flatten(0, 1)    # 展平输出维度
)

# 定义损失函数:均方误差(MSE),reduction='sum' 表示求和而非平均
loss_fn = torch.nn.MSELoss(reduction='sum')

learning_rate = 1e-6  # 学习率
for t in range(2000):  # 迭代2000次
    # 前向传播:直接调用模型,输入特征张量 xx,输出预测值 y_pred
    # Module 重载了 __call__ 方法,可像函数一样调用
    y_pred = model(xx)

    # 计算损失:传入预测值和真实值,返回损失张量
    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())  # 每100次迭代打印损失

    # 梯度清零:反向传播前必须清零,否则梯度会累积
    model.zero_grad()

    # 反向传播:自动计算模型所有可训练参数的梯度
    # 模型内部的参数默认设置 requires_grad=True
    loss.backward()

    # 梯度下降更新参数
    with torch.no_grad():
        # 遍历模型所有参数(weight 和 bias)
        for param in model.parameters():
            param -= learning_rate * param.grad

# 提取模型的第一层(线性层)
linear_layer = model[0]

# 线性层的参数:weight (1,3) 和 bias (1,)
# weight[:,0] → x 的系数,weight[:,1] → x² 的系数,weight[:,2] → x³ 的系数
print(f'Result: y = {linear_layer.bias.item()} + {linear_layer.weight[:, 0].item()} x + {linear_layer.weight[:, 1].item()} x^2 + {linear_layer.weight[:, 2].item()} x^3')

脚本运行耗时

0 分 0.537 秒

核心知识点解析

关键 API作用说明
torch.nn.Sequential按顺序堆叠神经网络层,构建模型
torch.nn.Linear(in_features, out_features)线性层(全连接层),实现 (y = xW + b)
torch.nn.Flatten(start_dim, end_dim)展平张量,合并指定维度
torch.nn.MSELoss(reduction)均方误差损失函数
- reduction='sum':损失求和
- reduction='mean':损失求平均(默认)
model.zero_grad()清零模型所有参数的梯度
model.parameters()返回模型所有可训练参数的迭代器
x.unsqueeze(dim)增加维度,-1 表示最后一维

总结

  1. nn 包是 PyTorch 构建神经网络的高层 API,nn.Sequential 可快速堆叠模块化的网络层;
  2. 线性层 nn.Linear 自动管理权重(weight)和偏置(bias),无需手动定义参数;
  3. 核心训练流程:前向传播 → 计算损失 → 梯度清零 → 反向传播 → 参数更新;
  4. 特征工程关键:将三阶多项式拟合转化为线性层问题,通过构造 (x, x², x³) 特征实现。