4. PyTorch深入：动态图机制、torch.nn模块、torch.utils.data数据加载、torch.distributed分布式训练

4.1 动态图机制：Define-by-Run 的核心原理

PyTorch 区别于 TensorFlow 1.x 静态图的核心在于其动态图机制（Define-by-Run）。计算图在每次前向传播时被实时构建，而非预先定义后再执行。这意味着 Python 控制流（if、for）可以直接参与图构建，使模型结构具备高度灵活性。

动态图的核心组件是 torch.Tensor 与 torch.autograd.Function。每个张量携带 requires_grad 标志和 grad_fn 指针，后者指向创建该张量的运算节点。反向传播时，autograd 引擎从损失张量出发，沿 grad_fn 链反向遍历，利用链式法则逐节点计算梯度。

import torch

x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

# 前向：动态构建计算图
y = w * x + b
z = torch.sigmoid(y)

# 反向：自动沿图回溯
z.backward()
print(f"dz/dw = {w.grad.item():.4f}")  # 输出约 0.0900
print(f"dz/db = {b.grad.item():.4f}")  # 输出约 0.0900

动态图机制带来三大优势：调试直观（可在任意中间节点打印张量）、结构自适应（如树形 LSTM、递归网络）、易于实现变长输入。代价是牺牲部分极致性能——每次迭代需重建图，且无法进行全局图优化。对于生产级大规模推理，可借助 torch.compile 或 torch.export 将动态图捕获为静态图以提升效率。

4.2 torch.nn 模块：从 Layer 到 Model 的构建体系

torch.nn 是构建神经网络的标准组件库，其核心抽象为 nn.Module。所有网络层（nn.Linear、nn.Conv2d、nn.LSTM）和完整模型均继承自该类。nn.Module 提供参数管理、子模块注册、设备迁移（.to(device)）、训练/评估模式切换（.train()/.eval()）等基础设施。

关键设计模式：在 __init__ 中声明层，在 forward 中定义计算逻辑。PyTorch 通过 __setattr__ 钩子自动收集 nn.Parameter 和子 nn.Module，因此 .parameters() 方法能递归返回所有可训练参数。

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, out_dim)

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

model = MLP(784, 256, 10)
print(f"可训练参数: {sum(p.numel() for p in model.parameters()):,}")  # 203,530

常用模块分类：基础层（Linear、Conv、Pooling、Dropout）、激活函数（ReLU、GELU、Softmax）、归一化（BatchNorm、LayerNorm）、循环层（RNN、LSTM、GRU）、损失函数（CrossEntropyLoss、MSELoss）。此外，nn.Sequential 适合线性堆叠结构，nn.ModuleList 和 nn.ModuleDict 用于管理动态数量的子模块。

4.3 torch.utils.data：高效数据加载流水线

数据加载是训练流程的起点，torch.utils.data 提供 Dataset 与 DataLoader 两大抽象。自定义数据集需继承 Dataset 并实现 __len__ 与 __getitem__ 方法。DataLoader 负责将数据集封装为可迭代批次，支持多进程并行加载、自动打乱、采样策略等。

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

class ImageDataset(Dataset):
    def __init__(self, paths, labels, transform=None):
        self.paths = paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        image = load_image(self.paths[idx])  # 自定义读取函数
        if self.transform:
            image = self.transform(image)
        return image, self.labels[idx]

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

dataset = ImageDataset(train_paths, train_labels, transform)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True,
                        num_workers=4, pin_memory=True)

关键参数调优要点：num_workers 通常设为 CPU 核心数的 2~4 倍；pin_memory=True 可加速 GPU 传输；drop_last=True 避免最后批次过小影响 BatchNorm 统计。对于超大数据集，可使用 IterableDataset 实现流式读取，或借助 torchdata 库构建可组合的数据管道。

4.4 torch.distributed：分布式训练核心机制

当单卡显存不足或训练时间过长时，需引入分布式训练。torch.distributed 提供进程组通信原语，支持数据并行（DDP）、模型并行和流水线并行。最常用的是 DistributedDataParallel，它通过 Ring-AllReduce 算法同步梯度，相比 DataParallel 具有更低的通信开销。

DDP 的核心流程：每个进程持有模型完整副本，处理不同数据分片；前向独立计算，反向时梯度通过 AllReduce 求平均并同步到所有进程。关键初始化步骤包括设置 init_method（常用 env://）、rank 和 world_size。

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def train_worker(rank, world_size):
    # 初始化进程组
    dist.init_process_group(backend='nccl', init_method='env://',
                            rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    model = MLP(784, 256, 10).to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # 使用 DistributedSampler 切分数据
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

    for epoch in range(10):
        sampler.set_epoch(epoch)  # 保证每个 epoch 打乱方式不同
        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.to(rank), batch_y.to(rank)
            optimizer.zero_grad()
            loss = criterion(ddp_model(batch_x), batch_y)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train_worker, args=(world_size,), nprocs=world_size)

启动命令通常使用 torchrun：torchrun --nproc_per_node=4 train.py。除 DDP 外，torch.distributed 还提供 all_gather、broadcast、reduce 等集合通信原语，用于实现自定义并行策略。对于超大模型（如 70B 参数），可结合 FullyShardedDataParallel 将参数、梯度和优化器状态分片到多卡。

4.5 案例：完整训练流程整合

以下展示一个将上述模块整合的典型训练脚本骨架，涵盖数据加载、模型构建、分布式封装与训练循环：

def main_worker(rank, world_size, config):
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    # 1. 数据
    train_ds = create_dataset(config.data_dir, train=True)
    sampler = DistributedSampler(train_ds, world_size, rank, shuffle=True)
    loader = DataLoader(train_ds, batch_size=config.batch_size,
                        sampler=sampler, num_workers=8, pin_memory=True)

    # 2. 模型
    model = create_model(config.model_name).to(rank)
    model = DDP(model, device_ids=[rank])

    # 3. 优化器与损失
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
    criterion = nn.CrossEntropyLoss()

    # 4. 训练循环
    for epoch in range(config.epochs):
        sampler.set_epoch(epoch)
        model.train()
        for images, labels in loader:
            images, labels = images.to(rank), labels.to(rank)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        # 5. 同步评估（仅 rank 0 打印）
        if rank == 0:
            acc = evaluate(model, val_loader)
            print(f"Epoch {epoch}: loss={loss.item():.4f}, acc={acc:.2f}%")

    dist.destroy_process_group()

4.6 常见误区与复盘

误区	后果	正确做法
在 forward 中创建新的 nn.Module	参数未注册，无法被优化器更新	所有层在 __init__ 中声明
DataLoader 设置 shuffle=True 且使用 DDP	每个 epoch 各进程数据划分不一致	使用 DistributedSampler 并调用 set_epoch()
评估时忘记调用 model.eval()	Dropout/BatchNorm 行为异常，指标失真	使用 with torch.no_grad() 且切换 eval 模式
DDP 中未将输入数据移至对应 GPU	设备不匹配报错或隐式拷贝降低性能	使用 batch.to(rank) 显式迁移
梯度累积时未手动清零	梯度跨 batch 累加，优化方向错误	每个 step 调用 optimizer.zero_grad()

复盘要点：动态图机制虽灵活，但需注意其内存开销——每次反向传播后计算图被释放，若需保留梯度用于二阶优化，应调用 retain_graph=True。在分布式场景中，通信与计算重叠是性能关键，DDP 已自动实现梯度 AllReduce 与反向传播的重叠，但需确保各卡计算负载均衡，避免掉队者拖慢整体训练速度。