Chuoer47's Blog

Appearance

神经网络模块 ​

torch.nn 是 PyTorch 中构建神经网络的核心模块。它提供了丰富的层（Layer）、损失函数、激活函数以及各种工具，使得构建、训练和部署深度学习模型变得简洁而高效。本文将全面讲解 nn.Module 的设计哲学以及各种常用网络层的使用方法。

一、nn.Module 基类详解 ​

nn.Module 是所有神经网络模块的基类。自定义模型必须继承它，并实现 __init__ 和 forward 方法。

1、基本结构 ​

python
import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()  # 必须调用父类构造函数
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 创建模型
model = SimpleNet(784, 256, 10)
print(model)

2、init 中的注意事项 ​

python
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()

        # 正确：将子模块赋值为属性
        self.linear = nn.Linear(10, 5)
        self.conv = nn.Conv2d(3, 16, 3)

        # 正确：使用 nn.ModuleList
        self.layers = nn.ModuleList([nn.Linear(10, 5) for _ in range(3)])

        # 错误：使用普通 list，参数不会被注册！
        # self.bad_layers = [nn.Linear(10, 5) for _ in range(3)]

        # 正确：注册非 nn.Module 的参数
        self.my_param = nn.Parameter(torch.randn(5, 5))

        # 正确：注册普通缓冲区（不参与优化，但会随模型保存）
        self.register_buffer("running_mean", torch.zeros(5))

只有赋值为属性的子模块才会被注册

如果子模块存储在普通 Python list 或 dict 中，它的参数不会出现在 model.parameters() 中，也就不会被优化器更新。务必使用 nn.ModuleList 或 nn.ModuleDict。

3、forward 方法 ​

python
class FlexibleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(32 * 8 * 8, 10)

    def forward(self, x):
        # forward 中可以使用任意 Python 控制流
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x

model = FlexibleNet()
# 注意：不要直接调用 model.forward(x)
# 应该使用 model(x)，它会自动处理钩子等逻辑
x = torch.randn(1, 3, 32, 32)
output = model(x)

为什么不该直接调用 forward()？

直接调用 model.forward(x) 会跳过 nn.Module.__call__ 中注册的钩子（hook）以及一些内部逻辑。始终使用 model(x) 的方式调用。

二、常用层详解 ​

1、全连接层（nn.Linear） ​

全连接层是最基础的层，执行线性变换 y=xWT+b。

python
import torch
import torch.nn as nn

# 参数：(输入特征数, 输出特征数, 是否使用偏置)
linear = nn.Linear(in_features=784, out_features=256, bias=True)

x = torch.randn(32, 784)  # (batch_size, in_features)
y = linear(x)
print(y.shape)  # torch.Size([32, 256])

# 查看参数
print(linear.weight.shape)  # torch.Size([256, 784])
print(linear.bias.shape)    # torch.Size([256])
参数类型说明
in_featuresint每个输入样本的特征数
out_featuresint每个输出样本的特征数
biasbool是否添加偏置项，默认 True

2、卷积层 ​

nn.Conv1d（一维卷积） ​

常用于序列数据（文本、时间序列、音频）。

python
import torch
import torch.nn as nn

# 参数：(输入通道, 输出通道, 卷积核大小)
conv1d = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3)

# 输入形状: (batch, channels, length)
x = torch.randn(32, 1, 100)  # batch=32, 1个通道, 序列长度=100
y = conv1d(x)
print(y.shape)  # torch.Size([32, 16, 98])  长度 = (100 - 3) / 1 + 1 = 98

nn.Conv2d（二维卷积） ​

最常用的卷积层，用于图像处理。

python
import torch
import torch.nn as nn

conv2d = nn.Conv2d(
    in_channels=3,      # 输入通道数（如 RGB 图像为 3）
    out_channels=64,    # 输出通道数（卷积核个数）
    kernel_size=3,      # 卷积核大小（3x3）
    stride=1,           # 步长
    padding=1,          # 填充
    dilation=1,         # 膨胀率
    groups=1,           # 分组卷积
    bias=True           # 是否使用偏置
)

# 输入形状: (batch, channels, height, width)
x = torch.randn(32, 3, 224, 224)
y = conv2d(x)
print(y.shape)  # torch.Size([32, 64, 224, 224])

# 输出尺寸计算公式：
# H_out = (H_in + 2*padding - dilation*(kernel_size-1) - 1) / stride + 1

nn.Conv3d（三维卷积） ​

用于视频、体积数据（如医学影像）。

python
conv3d = nn.Conv3d(in_channels=1, out_channels=32, kernel_size=3, padding=1)

# 输入形状: (batch, channels, depth, height, width)
x = torch.randn(8, 1, 16, 64, 64)  # 8个样本，16帧，64x64 分辨率
y = conv3d(x)
print(y.shape)  # torch.Size([8, 32, 16, 64, 64])

卷积参数速查 ​

参数说明常用值
kernel_size卷积核尺寸1, 3, 5, 7
stride滑动步长1（默认），2（下采样）
padding边缘填充0, 1, "same"
dilation膨胀卷积的空洞大小1（默认），2
groups分组卷积1（标准），in_channels（深度可分离）
bias是否加偏置True/False
python
# 深度可分离卷积（MobileNet 中的核心组件）
depthwise = nn.Conv2d(64, 64, kernel_size=3, padding=1, groups=64)
pointwise = nn.Conv2d(64, 128, kernel_size=1)

x = torch.randn(1, 64, 32, 32)
x = depthwise(x)   # 空间卷积
x = pointwise(x)   # 通道变换
print(x.shape)      # torch.Size([1, 128, 32, 32])

3、池化层 ​

最大池化与平均池化 ​

python
import torch
import torch.nn as nn

x = torch.randn(1, 1, 4, 4)
print(x)

# 最大池化
max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
y = max_pool(x)
print(y.shape)  # torch.Size([1, 1, 2, 2])

# 平均池化
avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
y = avg_pool(x)
print(y.shape)  # torch.Size([1, 1, 2, 2])

# 最大池化，返回索引（可用于 MaxUnpool2d）
max_pool_idx = nn.MaxPool2d(2, stride=2, return_indices=True)
y, indices = max_pool_idx(x)
print(y.shape)        # torch.Size([1, 1, 2, 2])
print(indices.shape)  # torch.Size([1, 1, 2, 2])

自适应池化（Adaptive Pool） ​

自适应池化允许指定输出尺寸，自动计算 kernel_size 和 stride，特别适合处理不同尺寸的输入。

python
import torch
import torch.nn as nn

# 自适应平均池化到指定输出大小
adaptive_pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))

x = torch.randn(1, 512, 7, 7)
y = adaptive_pool(x)
print(y.shape)  # torch.Size([1, 512, 1, 1])  → 全局平均池化

x2 = torch.randn(1, 512, 13, 13)
y2 = adaptive_pool(x2)
print(y2.shape)  # torch.Size([1, 512, 1, 1])  → 同样输出 1x1

# 自适应最大池化
adaptive_max = nn.AdaptiveMaxPool2d(output_size=(4, 4))
y3 = adaptive_max(torch.randn(1, 256, 50, 50))
print(y3.shape)  # torch.Size([1, 256, 4, 4])
池化层说明典型用途
nn.MaxPool2d取窗口内最大值CNN 下采样
nn.AvgPool2d取窗口内平均值CNN 下采样
nn.AdaptiveAvgPool2d指定输出大小的平均池化全局平均池化（GAP）
nn.AdaptiveMaxPool2d指定输出大小的最大池化特征提取
nn.GlobalAvgPool2d全局平均池化（等价于 AdaptiveAvgPool2d(1)）分类网络末端
nn.LPPool2dLp 范数池化特殊需求

4、归一化层 ​

BatchNorm（批归一化） ​

python
import torch
import torch.nn as nn

# 1D（全连接层后）
bn1d = nn.BatchNorm1d(num_features=64)
x = torch.randn(32, 64)  # (batch, features)
y = bn1d(x)
print(y.shape)  # torch.Size([32, 64])

# 2D（卷积层后，最常用）
bn2d = nn.BatchNorm2d(num_features=64)
x = torch.randn(32, 64, 28, 28)  # (batch, channels, H, W)
y = bn2d(x)
print(y.shape)  # torch.Size([32, 64, 28, 28])

# 3D
bn3d = nn.BatchNorm3d(num_features=32)
x = torch.randn(8, 32, 16, 16, 16)
y = bn3d(x)

# 关键参数
print(bn2d.weight.shape)   # 可学习的缩放参数 gamma
print(bn2d.bias.shape)     # 可学习的偏移参数 beta
print(bn2d.running_mean.shape)  # 运行均值（训练时更新，推理时使用）
print(bn2d.running_var.shape)   # 运行方差

LayerNorm（层归一化） ​

Transformer 中的标准归一化方式，不依赖 batch 维度。

python
import torch
import torch.nn as nn

# 对最后一个维度做归一化
ln = nn.LayerNorm(normalized_shape=512)

x = torch.randn(32, 10, 512)  # (batch, seq_len, d_model)
y = ln(x)
print(y.shape)  # torch.Size([32, 10, 512])

# 也可以对最后多个维度归一化
ln2 = nn.LayerNorm(normalized_shape=[10, 512])
y2 = ln2(x)
print(y2.shape)  # torch.Size([32, 10, 512])

GroupNorm（组归一化） ​

将通道分成若干组，每组内做归一化。不依赖 batch size，适用于小 batch 训练场景。

python
import torch
import torch.nn as nn

# 参数：(组数, 通道数)
gn = nn.GroupNorm(num_groups=32, num_channels=64)  # 64 通道分成 32 组，每组 2 通道

x = torch.randn(16, 64, 28, 28)
y = gn(x)
print(y.shape)  # torch.Size([16, 64, 28, 28])

# num_groups=1 时等价于 LayerNorm
# num_groups=num_channels 时等价于 InstanceNorm

InstanceNorm（实例归一化） ​

对每个样本的每个通道独立做归一化，常用于风格迁移。

python
import torch
import torch.nn as nn

inst_norm = nn.InstanceNorm2d(num_features=64)

x = torch.randn(16, 64, 28, 28)
y = inst_norm(x)
print(y.shape)  # torch.Size([16, 64, 28, 28])

归一化层对比 ​

归一化方式归一化维度依赖 Batch典型应用
BatchNorm(N, H, W) / (N, L)是CNN（训练时）
LayerNorm(C, H, W) / (L, D)否Transformer、RNN
GroupNorm(C//G, H, W)否小 batch CNN、Diffusion
InstanceNorm(H, W)否风格迁移

5、循环层 ​

RNN（基础循环神经网络） ​

python
import torch
import torch.nn as nn

# 参数：(输入特征数, 隐藏层大小, 层数)
rnn = nn.RNN(input_size=10, hidden_size=20, num_layers=2, batch_first=True)

# 输入: (batch, seq_len, input_size)
x = torch.randn(32, 50, 10)
output, h_n = rnn(x)

print(output.shape)  # torch.Size([32, 50, 20])  每个时间步的输出
print(h_n.shape)     # torch.Size([2, 32, 20])    最终隐藏状态

LSTM（长短期记忆网络） ​

python
import torch
import torch.nn as nn

lstm = nn.LSTM(
    input_size=10,
    hidden_size=20,
    num_layers=2,
    batch_first=True,
    bidirectional=True,  # 双向 LSTM
    dropout=0.1          # 层间 dropout
)

x = torch.randn(32, 50, 10)
output, (h_n, c_n) = lstm(x)

# 双向 LSTM 输出维度翻倍
print(output.shape)  # torch.Size([32, 50, 40])  20*2 = 40
print(h_n.shape)     # torch.Size([4, 32, 20])    2层*2方向 = 4
print(c_n.shape)     # torch.Size([4, 32, 20])    细胞状态

GRU（门控循环单元） ​

python
import torch
import torch.nn as nn

gru = nn.GRU(input_size=10, hidden_size=20, num_layers=2, batch_first=True)

x = torch.randn(32, 50, 10)
output, h_n = gru(x)

print(output.shape)  # torch.Size([32, 50, 20])
print(h_n.shape)     # torch.Size([2, 32, 20])
循环层门机制参数量说明
nn.RNN无最少最简单，容易梯度消失
nn.GRU重置门 + 更新门中等性价比高
nn.LSTM遗忘门 + 输入门 + 输出门最多表达能力最强

6、Dropout 层 ​

python
import torch
import torch.nn as nn

dropout = nn.Dropout(p=0.5)  # 50% 的概率将元素置零

x = torch.ones(10)

# 训练模式：随机丢弃
dropout.train()
y = dropout(x)
print(y)  # 一部分元素被置零，其余被放大 1/(1-p) 倍

# 推理模式：不做任何操作
dropout.eval()
y = dropout(x)
print(y)  # tensor([1., 1., 1., ...])

# Dropout 变体
dropout2d = nn.Dropout2d(p=0.2)   # 按通道丢弃（常用于卷积后）
dropout3d = nn.Dropout3d(p=0.2)   # 3D 版本
alpha_dropout = nn.AlphaDropout(p=0.5)  # 适用于 SELU 激活函数

Dropout 工作原理

训练时，Dropout 以概率 p 随机将元素置零，并将剩余元素乘以 11−p（inverted dropout），以保持期望值不变。推理时不做任何操作。

7、Embedding 层 ​

将离散的整数索引映射为密集的向量表示，是 NLP 模型的基础组件。

python
import torch
import torch.nn as nn

# 参数：(词汇表大小, 嵌入维度)
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=256)

# 输入必须是 LongTensor
x = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])  # (batch, seq_len)
y = embedding(x)
print(y.shape)  # torch.Size([2, 4, 256])

# 查看嵌入矩阵
print(embedding.weight.shape)  # torch.Size([10000, 256])

# 使用预训练词向量
pretrained = torch.randn(10000, 256)
embedding = nn.Embedding.from_pretrained(pretrained, freeze=False)
# freeze=False 表示嵌入向量也会被训练更新

# 填充索引（该索引对应的向量始终为零）
embedding = nn.Embedding(1000, 64, padding_idx=0)
x = torch.tensor([0, 1, 2, 0])
print(embedding(x))  # 索引 0 的嵌入始终为零向量
参数说明
num_embeddings词汇表大小
embedding_dim嵌入向量维度
padding_idx填充索引，对应的嵌入不会被更新
max_norm若指定，嵌入向量会被归一化到该范数
freezefrom_pretrained 专用，是否冻结嵌入

三、激活函数详解 ​

1、常用激活函数 ​

python
import torch
import torch.nn as nn
import torch.nn.functional as F

x = torch.linspace(-3, 3, 7)
print(f"输入: {x}")

# ReLU：max(0, x)
relu = nn.ReLU()
print(f"ReLU:     {relu(x)}")

# LeakyReLU：max(0.01*x, x)
leaky_relu = nn.LeakyReLU(negative_slope=0.01)
print(f"LeakyReLU:{leaky_relu(x)}")

# Sigmoid：1 / (1 + exp(-x))
sigmoid = nn.Sigmoid()
print(f"Sigmoid:  {sigmoid(x)}")

# Tanh：(exp(x) - exp(-x)) / (exp(x) + exp(-x))
tanh = nn.Tanh()
print(f"Tanh:     {tanh(x)}")

2、现代激活函数 ​

python
import torch
import torch.nn as nn

x = torch.linspace(-3, 3, 7)

# GELU：x * Φ(x)，Transformer 的标准激活函数
gelu = nn.GELU()
print(f"GELU:     {gelu(x)}")

# SiLU / Swish：x * sigmoid(x)
silu = nn.SiLU()
print(f"SiLU:     {silu(x)}")

# Mish：x * tanh(softplus(x))
mish = nn.Mish()
print(f"Mish:     {mish(x)}")

# ELU：max(0, x) + min(0, alpha * (exp(x) - 1))
elu = nn.ELU(alpha=1.0)
print(f"ELU:      {elu(x)}")

# SELU：自归一化激活函数
selu = nn.SELU()
print(f"SELU:     {selu(x)}")

3、Softmax ​

python
import torch
import torch.nn as nn

softmax = nn.Softmax(dim=-1)  # 指定在哪个维度上做 softmax

# 分类任务的典型用法
logits = torch.tensor([[2.0, 1.0, 0.1],
                        [0.5, 2.5, 0.3]])
probs = softmax(logits)
print(probs)
# tensor([[0.6590, 0.2424, 0.0986],
#         [0.1667, 0.6341, 0.1992]])
print(probs.sum(dim=-1))  # tensor([1.0000, 1.0000])

# LogSoftmax（数值更稳定，常与 NLLLoss 配合）
log_softmax = nn.LogSoftmax(dim=-1)
log_probs = log_softmax(logits)
print(log_probs)

4、激活函数对比 ​

激活函数公式值域特点典型应用
ReLUmax(0, x)[0, +inf)简单高效，可能神经元死亡CNN
LeakyReLUmax(αx, x)(-inf, +inf)缓解死亡问题CNN
Sigmoid1/(1+e^(-x))(0, 1)梯度消失，输出可做概率二分类输出
Tanh(e^x-e^(-x))/(e^x+e^(-x))(-1, 1)零中心化RNN
GELUx*Φ(x)(-0.17, +inf)平滑非单调Transformer
SiLU/Swishx*sigmoid(x)(-0.28, +inf)平滑非单调EfficientNet
Mishx*tanh(softplus(x))(-0.31, +inf)平滑自正则化YOLO v4
Softmaxe^(xi)/Σe^(xj)(0, 1)输出概率分布多分类输出
python
# 激活函数在模型中的典型用法
import torch.nn as nn

# CNN 中常用 ReLU
cnn_block = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(inplace=True)  # inplace=True 节省内存
)

# Transformer 中常用 GELU
transformer_ffn = nn.Sequential(
    nn.Linear(512, 2048),
    nn.GELU(),
    nn.Dropout(0.1),
    nn.Linear(2048, 512)
)

inplace 操作

ReLU(inplace=True) 会直接在输入张量上修改，节省内存。但在某些需要保留原始输入的场景下（如残差连接），使用 inplace=False（默认值）更安全。

四、容器模块 ​

1、nn.Sequential ​

按顺序依次执行各子模块，是最简洁的模型构建方式。

python
import torch
import torch.nn as nn

# 方式一：直接传入模块
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

x = torch.randn(32, 784)
output = model(x)
print(output.shape)  # torch.Size([32, 10])

# 方式二：使用 OrderedDict 给每层命名
from collections import OrderedDict

model = nn.Sequential(OrderedDict([
    ("fc1", nn.Linear(784, 256)),
    ("relu1", nn.ReLU()),
    ("dropout", nn.Dropout(0.2)),
    ("fc2", nn.Linear(256, 128)),
    ("relu2", nn.ReLU()),
    ("fc3", nn.Linear(128, 10)),
]))

# 可以按名称访问
print(model.fc1.weight.shape)  # torch.Size([256, 784])

2、nn.ModuleList ​

类似 Python list，但会正确注册子模块的参数。

python
import torch
import torch.nn as nn

class MultiLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(input_size if i == 0 else hidden_size, hidden_size)
            for i in range(num_layers)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

model = MultiLayerNet(784, 256, 5)
print(len(model.layers))  # 5
# 所有参数都被正确注册
print(sum(p.numel() for p in model.parameters()))

3、nn.ModuleDict ​

类似 Python dict，但会正确注册子模块的参数。

python
import torch
import torch.nn as nn

class MultiHeadNet(nn.Module):
    def __init__(self, input_size, num_classes_dict):
        super().__init__()
        self.shared_backbone = nn.Sequential(
            nn.Linear(input_size, 256),
            nn.ReLU(),
        )
        # 多个分类头
        self.heads = nn.ModuleDict({
            name: nn.Linear(256, num_classes)
            for name, num_classes in num_classes_dict.items()
        })

    def forward(self, x, task):
        features = self.shared_backbone(x)
        return self.heads[task](features)

model = MultiHeadNet(784, {"cls": 10, "det": 20, "seg": 5})
x = torch.randn(32, 784)
out = model(x, "cls")
print(out.shape)  # torch.Size([32, 10])

4、nn.ParameterList 与 nn.ParameterDict ​

与 ModuleList/ModuleDict 类似，但存储的是 nn.Parameter 而非 nn.Module。

python
import torch
import torch.nn as nn

class WeightedSum(nn.Module):
    def __init__(self, num_inputs):
        super().__init__()
        # 存储可学习的权重参数
        self.weights = nn.ParameterList([
            nn.Parameter(torch.randn(1)) for _ in range(num_inputs)
        ])

    def forward(self, inputs):
        result = torch.zeros_like(inputs[0])
        for w, x in zip(self.weights, inputs):
            result = result + w * x
        return result

model = WeightedSum(3)
inputs = [torch.randn(32, 10) for _ in range(3)]
output = model(inputs)
print(output.shape)  # torch.Size([32, 10])

容器模块对比 ​

容器存储类型访问方式注册参数典型用途
SequentialModule按顺序 / 按索引是简单线性模型
ModuleListModule按索引是动态数量的层
ModuleDictModule按名称是多任务头、多分支
ParameterListParameter按索引是可学习的参数集合
ParameterDictParameter按名称是可学习的参数字典

五、参数管理 ​

1、parameters() 与 named_parameters() ​

python
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

# parameters() 返回所有参数的迭代器
for param in model.parameters():
    print(param.shape)
# torch.Size([20, 10])   —— 第一层权重
# torch.Size([20])       —— 第一层偏置
# torch.Size([5, 20])    —— 第二层权重
# torch.Size([5])        —— 第二层偏置

# named_parameters() 返回 (名称, 参数) 对
for name, param in model.named_parameters():
    print(f"{name:30s} {str(param.shape):20s} requires_grad={param.requires_grad}")
# 0.weight                       torch.Size([20, 10]) requires_grad=True
# 0.bias                         torch.Size([20])     requires_grad=True
# 2.weight                       torch.Size([5, 20])  requires_grad=True
# 2.bias                         torch.Size([5])      requires_grad=True

# 常用操作
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数量: {total_params:,}")       # 325
print(f"可训练参数量: {trainable_params:,}")  # 325

2、state_dict() 与 load_state_dict() ​

state_dict 是一个 Python 字典，将每一层的参数映射到其对应的张量。

python
import torch
import torch.nn as nn

model = nn.Linear(10, 5)

# 获取状态字典
state = model.state_dict()
for name, tensor in state.items():
    print(f"{name}: {tensor.shape}")
# weight: torch.Size([5, 10])
# bias: torch.Size([5])

# 保存模型
torch.save(model.state_dict(), "model.pth")

# 加载模型
new_model = nn.Linear(10, 5)
new_model.load_state_dict(torch.load("model.pth"))

# strict=False 允许部分加载（模型结构不完全匹配时）
new_model2 = nn.Sequential(nn.Linear(10, 5), nn.ReLU())
# 可以只加载匹配的键
new_model2[0].load_state_dict(torch.load("model.pth"))

3、冻结与解冻参数 ​

python
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# 冻结前两层（迁移学习常见操作）
for name, param in model.named_parameters():
    if name.startswith("0.") or name.startswith("2."):
        param.requires_grad = False

# 查看冻结状态
for name, param in model.named_parameters():
    print(f"{name:30s} requires_grad={param.requires_grad}")

# 优化器只更新未冻结的参数
optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=0.001
)

# 解冻
for param in model.parameters():
    param.requires_grad = True

六、模型初始化 ​

1、PyTorch 默认初始化 ​

不同层有不同的默认初始化策略：

python
import torch.nn as nn

linear = nn.Linear(10, 5)
print(f"Linear 权重范围: [{linear.weight.min():.4f}, {linear.weight.max():.4f}]")
# 默认使用 Kaiming Uniform 初始化

conv = nn.Conv2d(3, 16, 3)
print(f"Conv2d 权重范围: [{conv.weight.min():.4f}, {conv.weight.max():.4f}]")

2、nn.init 模块 ​

python
import torch
import torch.nn as nn

linear = nn.Linear(10, 5)

# Xavier（Glorot）均匀分布初始化
nn.init.xavier_uniform_(linear.weight)
# Xavier 正态分布初始化
nn.init.xavier_normal_(linear.weight)

# Kaiming（He）均匀分布初始化（适配 ReLU）
nn.init.kaiming_uniform_(linear.weight, nonlinearity="relu")
# Kaiming 正态分布初始化
nn.init.kaiming_normal_(linear.weight, nonlinearity="relu")

# 常量初始化
nn.init.constant_(linear.bias, 0.0)

# 全零 / 全一初始化
nn.init.zeros_(linear.bias)
nn.init.ones_(linear.weight)

# 正交初始化（RNN 常用）
nn.init.orthogonal_(linear.weight)

# 稀疏初始化
nn.init.sparse_(linear.weight, sparsity=0.9)

3、初始化方法选择指南 ​

初始化方法适用激活函数公式说明
Xavier UniformSigmoid, TanhU(-sqrt(6/(fan_in+fan_out)), sqrt(6/(fan_in+fan_out)))保持前向/反向方差一致
Xavier NormalSigmoid, TanhN(0, sqrt(2/(fan_in+fan_out)))正态版本
Kaiming UniformReLU 及其变体U(-sqrt(6/fan_in), sqrt(6/fan_in))考虑 ReLU 的方差影响
Kaiming NormalReLU 及其变体N(0, sqrt(2/fan_in))最常用
OrthogonalRNN正交矩阵减轻梯度消失/爆炸

fan_in 与 fan_out

  • fan_in：输入特征数（权重矩阵的列数）
  • fan_out：输出特征数（权重矩阵的行数）

4、自定义初始化 ​

python
import torch
import torch.nn as nn

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d)):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)

# 对整个模型应用初始化
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
    nn.BatchNorm1d(10)
)
model.apply(init_weights)
python
# 更细粒度的初始化
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.fc = nn.Linear(64, 10)

        # 只初始化特定层
        nn.init.kaiming_normal_(self.conv1.weight, mode="fan_out")
        nn.init.constant_(self.bn1.weight, 1.0)
        nn.init.constant_(self.bn1.bias, 0.0)
        nn.init.xavier_uniform_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0.0)

七、模型信息查看 ​

1、基础信息查看 ​

python
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(64, 128, 3, padding=1),
    nn.BatchNorm2d(128),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(128, 10)
)

# 打印模型结构
print(model)

# 统计参数量
def count_parameters(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable = total - trainable
    return {
        "total": total,
        "trainable": trainable,
        "non_trainable": non_trainable
    }

params = count_parameters(model)
print(f"总参数量:     {params['total']:>12,}")
print(f"可训练参数量:  {params['trainable']:>12,}")
print(f"不可训练参数量:{params['non_trainable']:>12,}")

2、使用 torchinfo（推荐） ​

bash
pip install torchinfo
python
import torch
import torch.nn as nn
from torchinfo import summary

model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(64, 128, 3, padding=1),
    nn.BatchNorm2d(128),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(128, 10)
)

# 打印详细的模型摘要
summary(model, input_size=(1, 3, 32, 32))
# 输出包含每层的输出形状、参数量、MACs 等信息

3、参数量与模型大小估算 ​

python
import torch
import torch.nn as nn

def model_size_mb(model):
    """估算模型参数占用的内存大小（MB）"""
    param_size = 0
    buffer_size = 0

    for param in model.parameters():
        param_size += param.nelement() * param.element_size()

    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()

    total_mb = (param_size + buffer_size) / 1024 / 1024
    return total_mb

model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(64, 128, 3, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(128, 10)
)

print(f"模型大小: {model_size_mb(model):.4f} MB")

4、查看子模块 ​

python
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

# named_modules() 返回所有子模块
for name, module in model.named_modules():
    print(f"{name:20s} -> {module.__class__.__name__}")
# ""                   -> Sequential
# "0"                  -> Linear
# "1"                  -> ReLU
# "2"                  -> Linear

# children() 返回直接子模块
for child in model.children():
    print(child.__class__.__name__)

# modules() 返回所有模块（包括自身）
for module in model.modules():
    print(type(module).__name__)

参考资料