PyTorch神经网络层全攻略：从线性层到Transformer，一文掌握torch.nn核心组件

在深度学习框架PyTorch中，torch.nn模块提供了构建神经网络所需的各种层（Layer）。本文将系统介绍PyTorch中常用的层类型及其功能，帮助你更好地理解和搭建神经网络模型。

常用地址

官方文档：https://docs.pytorch.org/docs/main/nn.html

快速参考

容器（Containers）

容器层用于组织和管理其他层，是构建网络的基础。

Sequential示例

import torch
import torch.nn as nn

# 使用Sequential快速构建网络
model = nn.Sequential(
    nn.Linear(784, 1000),
    nn.ReLU(),
    nn.Dropout(),
    nn.Linear(1000, 10)
)

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

线性层（Linear Layers）

线性层执行仿射变换，是全连接网络的核心组件。

Linear示例

x = torch.randn(10, 24)

# 将24维输入映射到48维输出
fc = nn.Linear(in_features=24, out_features=48)
y = fc(x)
print(y.shape)  # torch.Size([10, 48])

卷积层（Convolution Layers）

卷积层通过滑动卷积核提取局部特征，是计算机视觉任务的基础。

Conv2d示例

# 输入: batch=1, 通道=3, 高度=32, 宽度=32
x = torch.randn(1, 3, 32, 32)

# 3输入通道 -> 16输出通道, 3x3卷积核
conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
y = conv(x)
print(y.shape)  # torch.Size([1, 16, 32, 32])

池化层（Pooling Layers）

池化层用于降采样，减少特征图尺寸和参数数量。

MaxPool2d示例

x = torch.randn(1, 16, 32, 32)

# 2x2池化窗口，步幅2，将32x32降采样为16x16
pool = nn.MaxPool2d(kernel_size=2, stride=2)
y = pool(x)
print(y.shape)  # torch.Size([1, 16, 16, 16])

# 自适应池化：输出固定尺寸5x5
adaptive_pool = nn.AdaptiveAvgPool2d(output_size=(5, 5))
y = adaptive_pool(x)
print(y.shape)  # torch.Size([1, 16, 5, 5])

填充层（Padding Layers）

填充层在输入边界添加填充值，常用于保持特征图尺寸。

正则化层（Normalization Layers）

正则化层用于稳定训练、加速收敛。

BatchNorm2d示例

x = torch.randn(16, 64, 28, 28)  # batch=16, channels=64

bn = nn.BatchNorm2d(num_features=64)
y = bn(x)
# 输出尺寸不变，但分布被归一化
print(y.shape)  # torch.Size([16, 64, 28, 28])

Dropout层（Dropout Layers）

Dropout通过随机丢弃神经元防止过拟合。

Dropout示例

x = torch.randn(100, 256)
dropout = nn.Dropout(p=0.5)  # 50%的神经元随机置零

model.train()  # 训练模式
y_train = dropout(x)

model.eval()   # 评估模式，dropout不生效
y_eval = dropout(x)

循环层（Recurrent Layers）

循环层用于处理序列数据，如时间序列、文本等。

LSTM示例

# 输入: 序列长度10, batch=5, 特征维度20
x = torch.randn(10, 5, 20)

# 输入20维，隐藏状态64维，2层LSTM
lstm = nn.LSTM(input_size=20, hidden_size=64, num_layers=2)

output, (h_n, c_n) = lstm(x)
print(output.shape)  # torch.Size([10, 5, 64])

Transformer层

Transformer层基于注意力机制，是现代NLP和视觉模型的核心。

TransformerEncoderLayer示例

x = torch.randn(10, 32, 512)  # 序列长10, batch=32, 特征512

encoder_layer = nn.TransformerEncoderLayer(
    d_model=512,      # 特征维度
    nhead=8,          # 8个注意力头
    dim_feedforward=2048,
    dropout=0.1
)

encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
y = encoder(x)
print(y.shape)  # torch.Size([10, 32, 512])

稀疏层（Sparse Layers）

稀疏层主要用于处理离散数据，如词嵌入。

Embedding示例

# 10个词的词汇表，每个词映射到64维向量
embedding = nn.Embedding(num_embeddings=10, embedding_dim=64)

# 输入: 索引序列
x = torch.tensor([1, 2, 4, 5, 9])  # shape: (5)
y = embedding(x)
print(y.shape)  # torch.Size([5, 64])

视觉层（Vision Layers）

视觉层提供图像特有的操作，如上采样。

完整示例：构建CNN分类网络

下面是一个综合使用多种层的完整示例：

import torch
import torch.nn as nn

class CNNClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        # 特征提取部分
        self.features = nn.Sequential(
            # 卷积块1
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 32x32 -> 16x16
            
            # 卷积块2
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 16x16 -> 8x8
            
            # 卷积块3
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4))  # 自适应输出4x4
        )
        
        # 分类器部分
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 10)  # 10个类别
        )
    
    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

# 测试模型
model = CNNClassifier()
x = torch.randn(16, 3, 32, 32)  # 16张32x32的RGB图像
y = model(x)
print(f"输入形状: {x.shape}")
print(f"输出形状: {y.shape}")  # torch.Size([16, 10])

总结

PyTorch的torch.nn模块提供了丰富的层类型，涵盖从基础线性层到高级Transformer层的各种需求。理解每种层的功能是构建高效神经网络的关键。建议在实际使用中查阅PyTorch官方文档获取更详细的参数说明和最新特性。