PyTorch：从张量、训练步骤到模型部署

从一个训练步骤认识 PyTorch

PyTorch 用 Tensor 表达数据，用自动求导记录可微计算，再由模型与优化器组织训练。学习时先掌握数据形状、类型、设备和梯度关系，比背诵 dir(torch) 的全部输出更有效。

本文的可执行示例固定在 Python 3.11、PyTorch 2.8.0 的 CPU 环境；配套视觉示例使用 torchvision 0.23.0。它们是验证基础行为的小实验，不包含真实数据集训练、精度评估或 GPU 性能结论。

一次训练中的问题	主要入口	本站后续文章
数据是什么形状，是否共享内存	torch.Tensor、torch	张量、存储与视图
如何组成模型与损失函数	torch.nn、torch.nn.functional	神经网络模块
算子怎样处理尺寸与边界	torch.nn、张量运算	算子与池化
怎样取样、组成小批量	torch.utils.data	数据集与 DataLoader
图像预处理、视觉模型和框运算	torchvision	torchvision 模块指南
怎样保存、恢复与部署模型	state_dict、torch.onnx	模型与设备操作

下面完整运行一次二分类训练步骤。输入是 4 个二维样本，线性层输出每个样本的两个 logits；CrossEntropyLoss 接收这些未归一化分数和整数类别。不要先对 logits 做 softmax 再交给该损失。

import torch
from torch import nn

torch.manual_seed(7)
x = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = torch.tensor([0, 1, 1, 0], dtype=torch.long)
model = nn.Linear(2, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()

model.train()
before = model.weight.detach().clone()
optimizer.zero_grad(set_to_none=True)
logits = model(x)
assert logits.shape == (4, 2)
loss = criterion(logits, y)
assert torch.isfinite(loss)
loss.backward()
assert all(p.grad is not None for p in model.parameters())
optimizer.step()
assert not torch.equal(before, model.weight)

model.eval()
with torch.inference_mode():
    prediction = model(x).argmax(dim=1)
assert prediction.shape == y.shape
print("loss before update:", float(loss.detach()))
print("predictions after one step:", prediction.tolist())

backward() 计算并累积梯度，step() 才更新参数，清梯度应放在需要开始新一轮累积的位置。train() / eval() 改变 Dropout、BatchNorm 等模块的行为；它们不控制自动求导开关。此例的线性层在两种模式下行为相同，但仍保留完整流程。一次更新只能证明计算链路正常，不能证明模型已学会异或关系或达到某种准确率。自动求导基础、优化器

按职责查找模块

模块	适合解决的问题	使用边界
torch、Tensor	创建、索引、广播、逐元素计算、归约	先核对 shape、dtype、device；形状相同不代表存储独立
torch.nn / functional	注册参数与缓冲区、组合网络、调用无状态运算	Module 的注册关系会影响保存、设备迁移与参数遍历
torch.autograd	梯度传播、梯度检查	原地修改、离散选择和 NumPy 往返都可能改变可微关系
torch.optim / lr_scheduler	更新参数、调整学习率	恢复训练通常也要恢复优化器、调度器状态
torch.utils.data	Dataset、Sampler、DataLoader	不自动理解图片语义；预处理由数据集或变换定义
torch.linalg、fft、special	线性代数、傅里叶变换、特殊函数	按数值问题选择算法与误差检查
torch.distributions、random	概率分布、随机数状态	固定种子不保证跨版本、跨设备逐位相同
torch.cuda、backends、amp	设备能力、后端与混合精度	可导入接口不代表机器具备可用 GPU
torch.profiler、torch.utils.benchmark	定位耗时、测量运算	GPU 是异步执行；不能直接用一次主机计时推断性能
torch.distributed、multiprocessing	多进程与分布式执行	需要额外的进程通信、数据分片和错误处理
torch.jit、torch.export、torch.onnx	脚本/跟踪、图捕获与交换格式	各自有支持范围；模型文件后缀不能说明采用了哪个机制

原记录中的 _C、_VF、_jit_internal、各类 Storage 与 JIT 内部类型，是某次安装环境的模块/类枚举。下划线内部入口、内置对象和过渡接口不能直接视为稳定学习清单；需要扩展算子时再查 torch.library、torch.ops 等对应文档。PyTorch 2.8 总目录

原有 22 组函数清单如何使用

原清单混合了模块函数、Tensor 方法、内部入口和重复分组。下表保留这些学习主题，把查找方向整理成可核查的入口。

原分组	查找方向与代表例子
1 数学运算	torch.add、exp；创建张量另查 zeros、randn；等差序列优先 arange
2 张量操作	reshape、cat、stack、比较与索引；视图/复制语义单独确认
3 张量变换	flip、rot90、roll，注意所作用的维度
4 线性代数	matmul、einsum、torch.linalg.svd / qr / solve
5 卷积和池化	nn.Conv2d、nn.MaxPool2d 及 nn.functional 对应函数
6 归一化与激活	nn.BatchNorm2d、LayerNorm、ReLU；训练/推理状态不是同一问题
7 损失函数	nn 或 nn.functional，核对输入域、目标类型和 reduction
8 优化器与调度	torch.optim、torch.optim.lr_scheduler
9 随机性	manual_seed、Generator、RNG 状态保存与恢复
10 属性与转换	x.dtype、x.device、x.to(...)、x.float()
11 逻辑归约	all、any、logical_and 等，明确归约维度
12 统计聚合	mean、sum、std、var、nanmean；注意校正参数与 NaN 策略
13 调试与求导状态	is_grad_enabled、no_grad、inference_mode、autograd.detect_anomaly
14 量化	quantize_per_tensor 等；量化张量和伪量化训练分别理解
15 索引切片	index_select、select、narrow、index_put
16 杂项张量操作	拆回形状、视图、索引和逐元素计算；谨慎使用 as_strided
17 特殊函数	查 torch.special，如误差函数与伽马相关函数
18 窗口函数	hann_window、hamming_window、kaiser_window
19 填充与散射	x.fill_、x.copy_、x.scatter_ / scatter_add_
20 条件操作	where、masked_select、masked_fill
21 排序与选择	sort、argsort、topk、quantile
22 统计	与第 12 组重叠，统一查统计归约接口

其中 torch.float 是 dtype，转换张量用 x.float() 或 x.to(torch.float32)；fill_、copy_ 通常应按 Tensor 方法查找。此处验证的 2.8 环境中没有 torch.nanstd、torch.nanvar、torch.scatter_fill_、torch.scatter_copy_、torch.set_inference_mode 这些入口。处理 NaN 方差需要明确过滤、计数和分母策略，不能把名字直接替换成 std；控制推理上下文使用 torch.inference_mode()。Tensor 参考、线性代数、inference_mode

检查当前环境，而不是复制旧环境常量

import torch

print("torch:", torch.__version__)
print("CUDA used to build this package:", torch.version.cuda)
print("CUDA available at runtime:", torch.cuda.is_available())
print("cuDNN version:", torch.backends.cudnn.version())
print("default dtype:", torch.get_default_dtype())
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
    print("device capability:", torch.cuda.get_device_capability(0))
    print("package architecture list:", torch.cuda.get_arch_list())

generator = torch.Generator(device="cpu").manual_seed(123)
state = generator.get_state()
first = torch.rand(5, generator=generator)
generator.set_state(state)
second = torch.rand(5, generator=generator)
assert torch.equal(first, second)
print("same CPU generator state reproduces this draw")

旧记录中的 has_cuda=False、has_mkl=True 等值只反映当时环境，不能当作安装要求。历史文档中“计算能力至少 3.0”的概括也不能作为当前安装包的 GPU 支持保证：要同时核对 PyTorch/torchvision 版本组合、安装包构建、驱动、硬件与目标运行时。torch.version.cuda 表示该 PyTorch 包所对应的 CUDA 构建版本，不等于系统 nvcc 的版本；仅安装 CUDA Toolkit 或 cuDNN 也不会把 CPU 版 PyTorch 自动变为 GPU 版。官方历史安装组合、CUDA 语义、可复现性说明

原笔记保存的 ApacheCN 中文文档、w3cschool PyTorch 教程 可作为历史学习入口。遇到接口差异时，以本机版本和相应版本的官方文档为准；CUDA/cuDNN 的当前安装条件另查 NVIDIA cuDNN 安装文档。