MILA深度学习教程：PyTorch神经网络库(torch.nn)全面解析

2025-06-20 17:15:35作者：霍妲思

前言

在深度学习领域，PyTorch因其动态计算图和直观的API设计而广受欢迎。作为PyTorch的核心组件之一，torch.nn库提供了构建神经网络所需的各种模块和功能。本文将深入解析torch.nn库的核心组件及其使用方法，帮助读者掌握PyTorch中神经网络构建的关键技术。

1. torch.nn库概述

torch.nn是PyTorch中专门用于构建神经网络的模块化库，它提供了构建各种神经网络所需的"积木块"。这些组件可以分为以下几大类：

1.1 基础层类型

线性层：nn.Linear（全连接层）、nn.Bilinear（双线性层）
卷积层：包括1D、2D、3D卷积(nn.Conv1d/2d/3d)和转置卷积(nn.ConvTranspose2d)
循环网络层：nn.LSTM、nn.GRU等

1.2 正则化与特殊层

归一化层：如nn.BatchNorm2d（批归一化）
Dropout层：nn.Dropout、nn.Dropout2d
嵌入层：nn.Embedding（常用于NLP任务）

1.3 激活函数与池化

非线性激活：nn.Sigmoid、nn.Tanh、nn.ReLU、nn.LeakyReLU等
池化层：nn.MaxPool1d、nn.AveragePool2d等

1.4 损失函数

包括nn.MSELoss（均方误差）、nn.CrossEntropyLoss（交叉熵）、nn.NLLLoss（负对数似然）等

2. 核心组件详解与示例

2.1 线性层与激活函数

线性层是神经网络最基本的构建块，下面展示如何使用nn.Linear和nn.Bilinear：

import torch
import torch.nn as nn

# 创建输入数据
x = torch.randn(32, 10)  # 32个样本，每个样本10维特征
y = torch.randn(32, 30)  # 32个样本，每个样本30维特征

# 定义层
linear = nn.Linear(in_features=10, out_features=20)  # 10维输入，20维输出
bilinear = nn.Bilinear(in1_features=10, in2_features=30, out_features=50)

# 前向传播
output_linear = linear(x)
output_bilinear = bilinear(x, y)

2.2 卷积层与批归一化

卷积神经网络(CNN)是计算机视觉任务的基础，PyTorch提供了完整的CNN构建模块：

# 输入数据：10个样本，3通道，28x28图像
x = torch.randn(10, 3, 28, 28)

# 定义CNN组件
conv = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
bn = nn.BatchNorm2d(num_features=32)  # 批归一化
pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化

# 前向传播
output_conv = bn(conv(x))  # 卷积+批归一化
output_pool = pool(conv(x))  # 卷积+池化

2.3 循环神经网络与嵌入层

对于序列数据，PyTorch提供了强大的循环网络支持：

# 输入数据：5个样本，每个样本3个token的索引
inputs = [[1, 2, 3], [1, 0, 4], [1, 2, 4], [1, 4, 0], [1, 3, 3]]
x = torch.LongTensor(inputs)

# 定义网络组件
embedding = nn.Embedding(num_embeddings=5, embedding_dim=20)  # 嵌入层
drop = nn.Dropout(p=0.5)  # Dropout层
gru = nn.GRU(input_size=20, hidden_size=50, num_layers=2, bidirectional=True)

# 前向传播
emb = drop(embedding(x))  # 嵌入+Dropout
gru_h, gru_h_t = gru(emb)  # GRU处理

3. 函数式API：torch.nn.functional

除了面向对象的API，PyTorch还提供了函数式编程接口torch.nn.functional（通常导入为F），它允许更灵活的网络构建方式：

import torch.nn.functional as F

x = torch.randn(10, 3, 28, 28)
filters = torch.randn(32, 3, 3, 3)

# 使用函数式API构建卷积网络
conv_out = F.relu(F.dropout(F.conv2d(input=x, weight=filters, padding=1), p=0.5))

函数式API特别适合需要自定义网络流程的场景，如实现复杂的注意力机制或自定义的残差连接。

4. 参数初始化：torch.nn.init

良好的参数初始化对神经网络训练至关重要，PyTorch提供了多种初始化方法：

import torch.nn.init as init

conv_layer = nn.Conv2d(3, 32, kernel_size=3, padding=1)

# 使用Kaiming初始化卷积权重
for name, param in conv_layer.named_parameters():
    if name == 'weight':
        init.kaiming_normal_(param)

常用初始化方法包括：

init.xavier_uniform_：Xavier均匀初始化
init.xavier_normal_：Xavier正态初始化
init.kaiming_normal_：Kaiming/He初始化
init.orthogonal_：正交初始化

5. 优化器：torch.optim

torch.optim模块提供了各种优化算法的实现：

import torch.optim as optim

# 定义可训练参数
W1 = torch.randn(10, 20, requires_grad=True)
W2 = torch.randn(10, 20, requires_grad=True)

# 创建优化器
optimizer = optim.Adam([W1, W2], lr=0.001, betas=(0.9, 0.999))

# 学习率调度器
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)