当前位置：首页 > news >正文

PyTorch 模型构建

news 2025/9/24 13:14:13

搭模型

nn.Module：所有模型 / 层的 “妈妈类”，想自定义模型或层，继承它准没错，还能自动管理参数（不用自己记权重）。

nn.Sequential：按顺序装层的 “盒子”，比如 “展平层→全连接层→激活层”，不用写复杂的连接逻辑，适合简单模型。

nn.functional：放无参数的 “小工具”，比如 ReLU 激活、池化，在模型里直接调用就行

1. 简单模型：用 nn.Sequential

适合层按顺序排的场景（比如普通分类模型），一行代码串起所有层：

python运行

import torch.nn as nn
# 输入28×28的图片，输出10个类别
model = nn.Sequential(nn.Flatten(),  # 把28×28展成784维向量nn.Linear(784, 300),  # 全连接层：784→300nn.ReLU(),  # 激活函数nn.Linear(300, 10),  # 输出层：300→10nn.Softmax(dim=1)  # 转成类别概率
)

2. 灵活模型：继承 nn.Module

如果层连接复杂（比如有分支），就自己定义模型，核心写 2 部分：

__init__：定义要用的层；

forward：定义数据怎么在层里走。

python运行

class MyModel(nn.Module):def __init__(self):super().__init__()# 1. 定义层self.flatten = nn.Flatten()self.linear1 = nn.Linear(784, 300)self.linear2 = nn.Linear(300, 10)def forward(self, x):# 2. 定义数据路径x = self.flatten(x)x = self.linear1(x)x = nn.functional.relu(x)  # 用nn.functional的激活x = self.linear2(x)return xmodel = MyModel()  # 实例化模型

3. 复杂模型：Module + 容器组合

如果模型分多个 “子模块”（比如 ResNet 的残差块），用 nn.Sequential 把子模块装起来，代码更清晰：

python运行

class ResBlock(nn.Module):# 自定义残差块（子模块）def __init__(self):super().__init__()self.conv = nn.Sequential(  # 子模块里用Sequentialnn.Conv2d(64, 64, 3, padding=1),nn.ReLU())def forward(self, x):return x + self.conv(x)  # 残差连接：输入+卷积输出# 用子模块搭完整模型
class ResNetSimple(nn.Module):def __init__(self):super().__init__()self.start = nn.Conv2d(3, 64, 3)self.blocks = nn.Sequential(  # 装多个残差块ResBlock(),ResBlock())def forward(self, x):x = self.start(x)x = self.blocks(x)return x

python运行

import torch
import torch.optim as optim# 1. 装数据（假设已做好预处理）
train_loader = ...  # 自己定义的DataLoader# 2. 初始化模型、损失、优化器
model = MyModel()
criterion = nn.CrossEntropyLoss()  # 分类任务
optimizer = optim.Adam(model.parameters(), lr=1e-4)# 3. 训练循环
for epoch in range(5):  # 训练5轮for data, label in train_loader:# 前向：算预测output = model(data)# 算损失loss = criterion(output, label)# 反向：清梯度→算梯度→更参数optimizer.zero_grad()loss.backward()optimizer.step()print(f"第{epoch+1}轮，损失：{loss.item():.4f}")