PyTorch 核心功能
三、核心功能
| 模块 | 能力 | 关键配置项 |
|---|---|---|
| torch(张量) | 创建、索引、广播、矩阵运算、device 迁移 | dtype、device、是否需要梯度 |
| autograd | 自动求导、梯度累积、自定义 Function | requires_grad、detach、no_grad |
| nn | 层、损失函数、容器(Sequential/ModuleList) | 初始化方式、参数量 |
| optim | SGD/Adam/AdamW 等优化器、学习率调度 | lr、weight_decay、betas |
| data | Dataset/DataLoader、采样器、collate_fn | batch_size、num_workers、pin_memory、shuffle |
| AMP | 混合精度训练,省显存提速 | autocast 的 dtype、GradScaler |
| 分布式 | DDP、FSDP、进程组、集合通信 | backend(nccl/gloo)、rank、world_size |
| torch.compile | 图编译加速 | mode、backend、dynamic shapes |
| 序列化 | state_dict 保存/加载、safetensors 互通 | map_location、weights_only |
| 推理 | eval 模式、no_grad、ONNX 导出 | 是否冻结、导出算子集 |
训练脚本里最该先设对的是这几个开关:model.train() 与 model.eval() 的切换(影响 Dropout 与 BatchNorm 行为)、optimizer.zero_grad() 的调用位置、以及 pin_memory=True 配合 non_blocking=True 搬数据。这三个细节写错,损失曲线会以很难解释的方式异常。

