常用层和函数#

PyTorch 提供了丰富的神经网络层和工具函数。本章将介绍常用的层类型和函数,帮助你更好地构建模型。

卷积层#

卷积层模块速查#

模块说明
nn.Conv1d一维卷积层,用于序列数据
nn.Conv2d二维卷积层,用于图像处理
nn.Conv3d三维卷积层,用于视频或3D数据
nn.ConvTranspose1d一维转置卷积(反卷积),用于上采样
nn.ConvTranspose2d二维转置卷积,用于生成模型和上采样
nn.ConvTranspose3d三维转置卷积
nn.Unfold将滑动窗口展开为矩阵
nn.Fold将展开的矩阵折叠回特征图

示例#

import torch
import torch.nn as nn

# Conv1d
conv1d = nn.Conv1d(16, 32, kernel_size=3, stride=1, padding=1)
x = torch.randn(32, 16, 100)  # (batch, channels, length)
out = conv1d(x)  # (32, 32, 100)

# Conv2d 基础用法
conv2d = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
x = torch.randn(32, 3, 32, 32)
out = conv2d(x)  # (32, 64, 32, 32)

# Conv3d
conv3d = nn.Conv3d(1, 32, kernel_size=3)
out = conv3d(torch.randn(32, 1, 16, 16, 16))

# 转置卷积(上采样)
transpose_conv = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
out = transpose_conv(torch.randn(32, 64, 16, 16))  # (32, 32, 32, 32)

# Unfold:将滑动窗口展开为矩阵
unfold = nn.Unfold(kernel_size=3, padding=1)
x = torch.randn(32, 64, 32, 32)
out = unfold(x)  # (32, 64*3*3, 32*32)

池化层#

池化层模块速查#

模块说明
nn.MaxPool1d一维最大池化
nn.MaxPool2d二维最大池化
nn.MaxPool3d三维最大池化
nn.AvgPool1d一维平均池化
nn.AvgPool2d二维平均池化
nn.AvgPool3d三维平均池化
nn.AdaptiveMaxPool1d一维自适应最大池化,输出固定大小
nn.AdaptiveMaxPool2d二维自适应最大池化
nn.AdaptiveMaxPool3d三维自适应最大池化
nn.AdaptiveAvgPool1d一维自适应平均池化
nn.AdaptiveAvgPool2d二维自适应平均池化(全局平均池化)
nn.AdaptiveAvgPool3d三维自适应平均池化
nn.MaxUnpool1d/2d/3d最大池化的逆操作
nn.FractionalMaxPool2d分数最大池化,使用随机步长
nn.LPPool1d/2dLp 范数池化

示例#

import torch
import torch.nn as nn

# MaxPool1d(1d 典型)
maxpool1d = nn.MaxPool1d(kernel_size=2, stride=2)
x = torch.randn(32, 64, 100)
out = maxpool1d(x)  # (32, 64, 50)

# MaxPool2d、AvgPool2d
maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
avgpool = nn.AvgPool2d(kernel_size=2, stride=2)
x = torch.randn(32, 64, 32, 32)
out = maxpool(x)   # (32, 64, 16, 16)
out = avgpool(x)   # (32, 64, 16, 16)

# 全局平均池化
global_avg = nn.AdaptiveAvgPool2d((1, 1))
out = global_avg(x)  # (32, 64, 1, 1)

归一化层#

归一化层模块速查#

模块说明
nn.BatchNorm1d一维批归一化,用于全连接层
nn.BatchNorm2d二维批归一化,用于卷积层
nn.BatchNorm3d三维批归一化
nn.GroupNorm组归一化
nn.LayerNorm层归一化
nn.InstanceNorm1d/2d/3d实例归一化,常用于风格迁移
nn.LocalResponseNorm局部响应归一化(LRN),常用于 AlexNet

示例#

import torch
import torch.nn as nn

# BatchNorm2d
bn = nn.BatchNorm2d(64)
x = torch.randn(32, 64, 32, 32)
out = bn(x)

# LayerNorm
ln = nn.LayerNorm(128)
x = torch.randn(32, 128)
out = ln(x)

# GroupNorm
gn = nn.GroupNorm(8, 64)
x = torch.randn(32, 64, 32, 32)
out = gn(x)

# InstanceNorm2d
in_norm = nn.InstanceNorm2d(64)
out = in_norm(x)

# LocalResponseNorm(LRN)
lrn = nn.LocalResponseNorm(size=5)
out = lrn(x)

激活函数#

非线性激活函数速查#

模块说明公式
nn.ReLU修正线性单元$\text{ReLU}(x) = \max(0, x)$
nn.LeakyReLU泄漏 ReLU$\text{LeakyReLU}(x) = \max(0, x) + \alpha \min(0, x)$
nn.PReLU参数化 ReLU,斜率可学习$\text{PReLU}(x) = \max(0, x) + a \min(0, x)$
nn.RReLU随机 ReLU,训练时随机选择斜率-
nn.SigmoidSigmoid 函数$\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}$
nn.Tanh双曲正切函数$\text{Tanh}(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}$
nn.GELU高斯误差线性单元$\text{GELU}(x) = x \times \Phi(x)$
nn.SiLUSigmoid 线性单元(Swish)$\text{SiLU}(x) = x \times \text{sigmoid}(x)$
nn.ELU指数线性单元$\text{ELU}(x) = \max(0, x) + \min(0, \alpha \times (\exp(x) - 1))$
nn.CELU连续可微的 ELU-
nn.SELU缩放指数线性单元-
nn.ReLU6ReLU6$\text{ReLU6}(x) = \min(\max(0, x), 6)$
nn.HardswishHardswish 激活函数-
nn.HardsigmoidHard sigmoid 激活函数-
nn.HardtanhHard tanh 激活函数-
nn.SoftmaxSoftmax 函数$\text{Softmax}(x_i) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$
nn.LogSoftmax对数 Softmax$\text{LogSoftmax}(x_i) = \log\left(\frac{\exp(x_i)}{\sum_j \exp(x_j)}\right)$
nn.SoftminSoftmin 函数(Softmax 的负值版本)-
nn.Softmax2d对每个空间位置应用 Softmax-
nn.SoftsignSoftsign 函数$\text{Softsign}(x) = \frac{x}{1 + \|x\|}$
nn.TanhshrinkTanhshrink 函数$\text{Tanhshrink}(x) = x - \tanh(x)$
nn.Threshold阈值激活函数-

示例#

import torch
import torch.nn as nn

x = torch.randn(32, 128)

# ReLU、LeakyReLU
relu = nn.ReLU()
out = relu(x)
leaky_relu = nn.LeakyReLU(negative_slope=0.01)
out = leaky_relu(x)

# Sigmoid、Tanh、GELU
out = nn.Sigmoid()(x)
out = nn.Tanh()(x)
out = nn.GELU()(x)

# Softmax、ReLU6
out = nn.Softmax(dim=-1)(x)
out = nn.ReLU6()(x)

Dropout 层#

Dropout 层模块速查#

模块说明
nn.DropoutDropout,训练时随机置零
nn.Dropout1d/2d/3d多维 Dropout
nn.AlphaDropoutAlpha Dropout,保持自归一化性质
nn.FeatureAlphaDropout特征 Alpha Dropout

示例#

import torch
import torch.nn as nn

# Dropout(全连接)
dropout = nn.Dropout(p=0.5)
x = torch.randn(32, 128)
out = dropout(x)

# Dropout2d(卷积特征图)
dropout2d = nn.Dropout2d(p=0.5)
x = torch.randn(32, 64, 32, 32)
out = dropout2d(x)

全连接层#

模块说明公式
nn.Linear全连接层(线性变换)$y = xA^T + b$
nn.Bilinear双线性层$y = x_1^T A x_2 + b$

示例#

import torch
import torch.nn as nn

# Linear:会自动初始化 权重(Kaiming) 均匀分布;偏置(均匀分布)
fc = nn.Linear(784, 128) # nn.Linear(784,128):特征维度
x = torch.randn(32, 784) # x维度是:batch维 + 特征维
out = fc(x)  # (32, 128)  

# Bilinear
bilinear = nn.Bilinear(64, 64, 32)
x1 = torch.randn(32, 64)
x2 = torch.randn(32, 64)
out = bilinear(x1, x2)  # (32, 32)

# 多层全连接(MLP)
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

循环神经网络层#

循环层模块速查#

模块说明
nn.RNN基础循环神经网络
nn.LSTM长短期记忆网络
nn.GRU门控循环单元
nn.RNNCell单个时间步的 RNN 单元
nn.LSTMCell单个时间步的 LSTM 单元
nn.GRUCell单个时间步的 GRU 单元

示例#

import torch
import torch.nn as nn

x = torch.randn(32, 5, 10)  # (batch, seq_len, input_size)

# LSTM
lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=2, batch_first=True)
out, (h_n, c_n) = lstm(x)  # out: (32, 5, 20)

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

# LSTMCell(单步)
cell = nn.LSTMCell(10, 20)
h, c = torch.zeros(32, 20), torch.zeros(32, 20)
for t in range(5):
    h, c = cell(x[:, t], (h, c))

嵌入层#

稀疏层模块速查#

模块说明
nn.Embedding嵌入层,将索引映射为密集向量
nn.EmbeddingBag嵌入包,计算嵌入向量的和或均值

示例#

import torch
import torch.nn as nn

# Embedding
embedding = nn.Embedding(1000, 128)
x = torch.randint(0, 1000, (32, 10))
out = embedding(x)  # (32, 10, 128)

# EmbeddingBag(按包聚合,如词袋)
embed_bag = nn.EmbeddingBag(1000, 128, mode='mean')
offsets = torch.tensor([0, 4, 8, 12, 16, 20, 24, 28, 32])  # 8 个包
indices = torch.randint(0, 1000, (32,))
out = embed_bag(indices, offsets)  # (8, 128)

注意力机制#

Transformer 层模块速查#

模块说明
nn.Transformer完整的 Transformer 模型
nn.TransformerEncoderTransformer 编码器
nn.TransformerDecoderTransformer 解码器
nn.TransformerEncoderLayer单个 Transformer 编码器层
nn.TransformerDecoderLayerTransformer 解码器层
nn.MultiheadAttention多头注意力机制

示例#

import torch
import torch.nn as nn

x = torch.randn(32, 10, 512)  # (batch, seq_len, embed_dim)

# MultiheadAttention
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
out, attn_weights = mha(x, x, x)  # out: (32, 10, 512)

# TransformerEncoderLayer
enc_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
out = enc_layer(x)  # (32, 10, 512)

实用工具函数#

工具函数速查#

模块说明
nn.Flatten展平层
nn.Unflatten反展平层
nn.Identity恒等映射,不做任何操作

示例(Flatten / Unflatten / Identity)#

import torch
import torch.nn as nn

# Flatten
flatten = nn.Flatten()
x = torch.randn(32, 3, 32, 32)
out = flatten(x)  # (32, 3072)

# Unflatten
unflatten = nn.Unflatten(1, (3, 32, 32))
out = unflatten(torch.randn(32, 3072))  # (32, 3, 32, 32)

# Identity
identity = nn.Identity()
out = identity(x)  # 形状不变

填充层速查#

模块说明
nn.ZeroPad2d零填充
nn.ConstantPad1d/2d/3d常数填充
nn.ReflectionPad1d/2d/3d反射填充
nn.ReplicationPad1d/2d/3d复制填充

示例(Padding)#

import torch
import torch.nn as nn

x = torch.randn(32, 3, 32, 32)

# ZeroPad2d
pad = nn.ZeroPad2d(padding=1)
out = pad(x)  # (32, 3, 34, 34)

# ConstantPad2d
pad2 = nn.ConstantPad2d(padding=1, value=0)
out = pad2(x)

组合层#

容器类模块速查#

模块说明
nn.Module所有神经网络模块的基类
nn.Sequential按顺序组合多个模块
nn.ModuleList存储模块的列表,可动态添加
nn.ModuleDict存储模块的字典
nn.ParameterList存储参数的列表
nn.ParameterDict存储参数的字典

示例#

import torch
import torch.nn as nn

# Sequential
model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(64 * 16 * 16, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# ModuleList
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(10, 20),
            nn.Linear(20, 40)
        ])
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# ParameterList
class Custom(nn.Module):
    def __init__(self):
        super().__init__()
        self.params = nn.ParameterList([nn.Parameter(torch.randn(4, 4)) for _ in range(3)])
    def forward(self, x):
        for p in self.params:
            x = x @ p
        return x

常用函数#

损失函数速查#

模块说明公式
nn.L1LossL1 损失(平均绝对误差)$L = \frac{1}{N} \sum_i \|x_i - y_i\|$
nn.MSELoss均方误差损失$L = \frac{1}{N} \sum_i (x_i - y_i)^2$
nn.CrossEntropyLoss交叉熵损失,用于多分类-
nn.NLLLoss负对数似然损失,通常与 LogSoftmax 配合使用-
nn.BCELoss二元交叉熵损失,用于二分类-
nn.BCEWithLogitsLoss带 logits 的二元交叉熵,数值更稳定-
nn.KLDivLossKL 散度损失-
nn.MarginRankingLoss排序损失-
nn.HingeEmbeddingLossHinge 嵌入损失-
nn.MultiLabelMarginLoss多标签边界损失-
nn.SmoothL1Loss平滑 L1 损失(Huber 损失)-
nn.SoftMarginLossSoft margin 损失-
nn.MultiLabelSoftMarginLoss多标签 Soft margin 损失-
nn.CosineEmbeddingLoss余弦嵌入损失-
nn.MultiMarginLoss多类多分类损失-
nn.TripletMarginLoss三元组损失-
nn.CTCLoss连接时序分类损失,用于序列标注-
nn.PoissonNLLLoss泊松负对数似然损失-
nn.GaussianNLLLoss高斯负对数似然损失-

示例(F.* 与 nn 损失)#

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

# F.* 典型
x = torch.randn(32, 64, 32, 32)
x = F.relu(x)
x = F.max_pool2d(x, kernel_size=2, stride=2)
x = F.dropout(x, p=0.5, training=True)
x = F.interpolate(x, size=(64, 64), mode='bilinear', align_corners=False)

# F 损失
logits = torch.randn(32, 10)
targets = torch.randint(0, 10, (32,))
loss = F.cross_entropy(logits, targets)
loss = F.mse_loss(torch.randn(32, 1), torch.randn(32, 1))

# nn 损失
ce = nn.CrossEntropyLoss()
loss = ce(logits, targets)
bce = nn.BCEWithLogitsLoss()
loss = bce(torch.randn(32, 1), torch.rand(32, 1))

构建复杂模型示例#

ResNet Block#

import torch
import torch.nn as nn

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride=stride),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        residual = self.shortcut(x)
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += residual
        out = self.relu(out)
        return out

Transformer Block#

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, embed_dim),
            nn.Dropout(dropout)
        )
    
    def forward(self, x):
        # 自注意力
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)
        
        # 前馈网络
        ff_out = self.feed_forward(x)
        x = self.norm2(x + ff_out)
        
        return x

练习#

  1. 使用不同的层构建一个 CNN 模型
  2. 实现一个带残差连接的模块
  3. 构建一个简单的 Transformer 块
  4. 尝试使用不同的归一化层,比较效果

下一步#

掌握了常用层和函数后,学习:

其他模块参考#

以下模块未在上文展开,仅列出速查表;完整列表见 附录1-PyTorch 完整参考

距离函数(Distance Functions)#

模块说明
nn.CosineSimilarity余弦相似度
nn.PairwiseDistance成对距离(Lp 范数)
import torch
import torch.nn as nn

cos = nn.CosineSimilarity(dim=1)
x1, x2 = torch.randn(32, 64), torch.randn(32, 64)
out = cos(x1, x2)  # (32,)

视觉层(Vision Layers)#

模块说明
nn.PixelShuffle像素重排,用于上采样
nn.PixelUnshuffle像素反重排,PixelShuffle 的逆操作
nn.Upsample上采样层
nn.UpsamplingNearest2d最近邻上采样
nn.UpsamplingBilinear2d双线性上采样
import torch
import torch.nn as nn

ps = nn.PixelShuffle(upscale_factor=2)
x = torch.randn(32, 64, 16, 16)
out = ps(x)  # (32, 16, 32, 32)

数据并行层(DataParallel Layers)#

模块说明
nn.DataParallel数据并行包装器(已弃用,推荐使用 DistributedDataParallel)
nn.DistributedDataParallel分布式数据并行(推荐使用)

参数初始化(Initialization)#

torch.nn.init 模块提供参数初始化函数:

函数说明
init.xavier_uniform_Xavier/Glorot 均匀初始化
init.xavier_normal_Xavier/Glorot 正态初始化
init.kaiming_uniform_Kaiming/He 均匀初始化(适用于 ReLU)
init.kaiming_normal_Kaiming/He 正态初始化(适用于 ReLU)
init.uniform_均匀分布初始化
init.normal_正态分布初始化
init.constant_常数初始化
init.zeros_零初始化
init.ones_一初始化
init.orthogonal_正交初始化
init.sparse_稀疏初始化
init.trunc_normal_截断正态分布初始化
import torch
import torch.nn as nn
import torch.nn.init as init
m = nn.Linear(64, 32)
init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
init.zeros_(m.bias)