★ FEATURED ARTICLE PyTorch 口袋参考这个系列我打算写成一个可以随手翻开即查即用的东西。第一篇先从最核心的张量、自动求导和模型搭建讲起里面穿插了我平时调试时踩过的坑和总结出来的一些“肌肉记忆”写法。适合刚把 PyTorch 装好、准备动手写训练代码的读者也适合已经写过几个模型、但想把自己零散的思路重新规整一遍的朋友。我尽量不聊大理论而是把“到底怎么写才能跑通、跑稳”讲清楚代码片段都能直接复制到你的脚本里稍作改动就能用。1. 张量先把这个基础夯实1.1 创建张量的几种常用姿势PyTorch 的一切操作基本都围绕着张量Tensor展开它本质上就是一个可以放到 GPU 上计算的多维数组。创建张量的方式很多但你实际写代码时常用的其实就那么几个。最基础的是torch.tensor()它接收一个 Python 列表或 NumPy 数组import torch a torch.tensor([1, 2, 3]) # 整数张量dtype 默认是 torch.int64 b torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # 浮点张量dtype 默认是 torch.float32这里有个小细节容易忽略torch.tensor()会直接复制数据如果你从 NumPy 数组创建改变原数组不会影响这个张量。而torch.from_numpy()则是共享底层内存的修改任意一方另一方都会跟着变。用 NumPy 做数据预处理时我经常用后者能省一次拷贝但也要留意别在错误的地方改了共享数据。创建全 0、全 1、随机初始化张量时推荐用带_like的系列函数比如torch.zeros_like(x)、torch.ones_like(x)。它们会自动继承x的 dtype、device 和 shape写起来更省事也不会因为设备不一致报错。随机张量需要注意分布。torch.randn()生成标准正态分布适合作为初始化噪声torch.rand()生成从 0 到 1 均匀分布。初始化权重时业界常用nn.init.kaiming_uniform_()这类方法直接原地修改张量训练收敛明显更稳定。1.2 dtype、device 与 shape三个最容易翻车的属性张量有三个属性我每次排查 bug 都会先看一遍dtype数据类型、device所在设备和 shape形状。dtype 出问题最典型的表现是“loss 变成 nan 或直接报错”。PyTorch 默认浮点类型是float32但有些老模型或存储层拿到的是float64和默认权重一计算就可能报“expected scalar type Double but found Float”。统一用.float()或者.to(dtypetorch.float32)最省心。另外布尔张量不能直接求梯度整数张量也不能做需要梯度的运算想参与自动求导就必须转成浮点。device 不一致的报错信息通常长这样Expected all tensors to be on the same device。模型在 GPU 上输入在 CPU 上一 forward 就炸。我现在的习惯是数据进 model 之前统一执行batch {k: v.to(device) for k, v in batch.items()}模型只需要在创建后执行一次model.to(device)。shape 问题就更常见了。全连接层要求输入是二维(batch, features)卷积层要求四维(batch, channels, height, width)一不留神少一个维度就报错。调试时多用print(x.shape)或者直接把assert x.shape torch.Size([2, 3, 224, 224])写在关键节点前面能省大量排查时间。2. 索引、切片与变形2.1 view、reshape、permute 到底怎么选张量的变形操作很容易让人迷惑尤其是view、reshape、permute这三个。view是在原张量的内存上重新解释维度要求张量在内存中是连续存储的。如果一个张量是经过转置、切片得到的它很可能不再连续直接view会报错。解决办法是先用.contiguous()保证内存连续或者干脆用reshape。reshape在没有连续性的情况下会自动复制一份数据行为更宽容所以我建议在不确定的时候优先用reshape。但这两者都不能做维度交换。想交换维度需要用permute或者transpose。比如图像张量从(batch, channel, height, width)改成(batch, height, width, channel)一条x.permute(0, 2, 3, 1)就够了。transpose只能交换两个维度permute可以一次性重排所有维度。这里要特别提醒permute返回的张量和原张量共享内存修改它会改到原始数据。如果需要断开关联记得加.clone()。2.2 聚合操作里的 keepdim 陷阱求均值、求和是每个训练脚本都会用到的操作。torch.mean(x, dim1)会默认把指定维度“挤掉”导致结果少一个维度。比如输入是[4, 3, 224, 224]对 dim1 做 mean 后得到[4, 224, 224]。如果你之后还要做广播操作这个被挤掉的维度可能会让次方运算报错。解决办法是加keepdimTruetorch.mean(x, dim1, keepdimTrue)返回[4, 1, 224, 224]。这样在和原始形状做广播时会更安全代码意图也更清楚。还有个实用技巧求最大值位置用torch.argmax(x, dim1)得到的是 index 张量。如果后续需要把这个 index 变成 one-hot直接用torch.nn.functional.one_hot就行但注意它不接受(batch, seq)形状的张量直接转换需要先把维度展平或者设置num_classes。3. autograd自动求导没那么神秘3.1 计算图与 requires_grad 的关系PyTorch 的自动求导机制核心是动态计算图。每做一次张量运算它内部就会记录这次运算的输入、输出和梯度函数形成一个有向无环图。当你调用.backward()时它从当前输出出发沿反向路径把梯度逐层算回叶子节点。只有设置了requires_gradTrue的叶子张量才会真正累积梯度。模型的参数默认就是这种叶子张量。而中间计算结果默认requires_gradFalse但如果你对中间结果调用.backward()它反而会报错因为只有标量才能直接 backward。想让网络输出是一个向量也能 backward需要传入同样形状的gradient参数作为上游梯度。我在实践中很少手动设置张量的requires_grad因为nn.Parameter已经帮你处理好了。需要冻结某些网络层时更推荐用parameter.requires_grad_(False)这样它在推理时不会计算梯度省显存也省时间。3.2 backward 与梯度累积的坑一个常见的误解是每次.backward()之后参数的梯度是自动清零的。恰恰相反默认情况下梯度会累积到.grad里。如果你在循环里不断 backward 而忘记optimizer.zero_grad()梯度会越叠越大训练直接发散。标准流程是三步先optimizer.zero_grad()再loss.backward()最后optimizer.step()。顺序不要写反因为zero_grad()必须在前一步把旧梯度清掉。另外.backward()默认只计算叶子节点的梯度中间变量不会保留 grad。如果你想调试某一个中间层的梯度可以在 forward 时对那个张量做x.register_hook(lambda grad: print(grad))打印出不经过该张量后的梯度。这个技巧在排查梯度消失、梯度爆炸时非常管用。分布式训练或多卡场景下梯度同步也是很多人踩坑的点。用torch.nn.parallel.DistributedDataParallel后backward 会自动做 all-reduce 同步参数梯度前提是loss被所有卡以相同方式计算。如果你在每张卡上用了不同 batch size最后平均梯度时会出现差异这会影响收敛。4. 模型定义从 nn.Module 到实战4.1 用 nn.Sequential 快速堆叠用 nn.Module 实现复杂控制流最简单的模型定义当然是nn.Sequential它适合按顺序执行的网络import torch.nn as nn model nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) )缺点是它不擅长分支结构。如果你的 forward 里要同时走两条路径再把结果拼接或者需要根据条件执行不同层就应该用自定义nn.Moduleclass MyNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 256) self.fc2 nn.Linear(256, 10) def forward(self, x): x torch.relu(self.fc1(x)) return self.fc2(x)把层定义在__init__里把计算逻辑写在forward里这是最基本的规范。不要直接在forward里临时创建新的nn.Linear那样每次前向都会新建一组参数模型根本无法保存和复用。4.2 自定义层和 forward 的常见错误自定义层时最容易犯的错误是忘记调用super().__init__()。在nn.Module的__init__中如果没有执行super().__init__()后续注册参数和模块的机制就会失效model.parameters()返回空列表训练时 optimizer 什么都没更新。另一个常见错误是在forward里修改了输入张量的数量或形状却没有在__init__里对应调整。比如两分支相加时需要保证两个分支输出的维度一致。我通常会在分支前分别用nn.AdaptiveAvgPool2d((1, 1))把特征图压成固定尺寸再做torch.cat。使用nn.ModuleList维护多个子层也值得掌握。直接写self.layers [nn.Linear(...) for _ in range(3)]会创建一个普通列表里面的参数不会被注册。正确写法是self.layers nn.ModuleList([...])这样所有参数都能被parameters()访问到。5. 数据加载与训练循环5.1 Dataset 和 DataLoader 的正确用法数据加载是训练代码里最容易写乱的部分。一个标准的torch.utils.data.Dataset需要实现三个方法__init__、__len__和__getitem__。__getitem__返回一个样本通常是一个元组输入标签或者字典。DataLoader会在内部调用__getitem__来组装 batch。如果每个样本的尺寸不一样比如文本长度不同DataLoader默认会尝试把它们堆成张量形状不一致就会报错。这时你需要传递一个collate_fn函数手动把样本列表整理成一个 batchdef collate_fn(batch): inputs, labels zip(*batch) inputs torch.nn.utils.rnn.pad_sequence(inputs, batch_firstTrue, padding_value0) return inputs, torch.stack(labels)这里的pad_sequence可以将长度不同的序列填充到同一个长度。填充之后还要注意在损失函数里过滤填充部分否则模型会在 padding 上浪费时间。DataLoader有个参数num_workers多进程加载数据能显著提速但 Windows 下写成if __name__ __main__保护是必须的否则会无限递归创建子进程。Linux 下一般可以设成 CPU 核心数的一半但并不是越多越好过多的 worker 反而会拖慢速度。5.2 优化器、损失函数与一个完整的训练骨架PyTorch 里最常见的优化器是torch.optim.Adam和torch.optim.SGD。我自己的习惯是图像分类任务先用 Adam 快速找到好的初始点再切到 SGD momentum 做微调NLP 任务基本一直用 AdamW配合 warmup 效果更稳。损失函数的选择要跟任务匹配。分类用nn.CrossEntropyLoss它内部已经把 softmax 和 log 都处理了所以网络最后一层不要手动加 softmax。回归用nn.MSELoss或nn.L1Loss。多标签分类用nn.BCEWithLogitsLoss同样不推荐手动加 sigmoid。训练骨架我一般是这样写的model MyNet().to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() for epoch in range(num_epochs): model.train() total_loss 0 for x, y in dataloader: x, y x.to(device), y.to(device) optimizer.zero_grad() out model(x) loss criterion(out, y) loss.backward() optimizer.step() total_loss loss.item() # 验证 model.eval() with torch.no_grad(): correct 0 for x, y in val_loader: x, y x.to(device), y.to(device) out model(x) pred out.argmax(dim1) correct (pred y).sum().item()注意model.train()和model.eval()之间的切换。Dropout和BatchNorm在两种模式下的行为完全不同。不加model.eval()你的验证集指标往往会偏高因为 dropout 在验证时没有被关闭。还有一个细节是loss.item()。如果你直接把loss累加它会带着计算图导致显存不断累积。.item()能把它转成 Python 浮点数。6. 版本、设备与调试6.1 安装 PyTorch 时怎么选版本PyTorch 的安装直接决定你后面踩多少坑。官方推荐的 pip 命令通常类似pip install torch torchvision torchaudio但这条命令在你的机器上可能装成 CPU 版本也可能装成和 CUDA 版本不匹配的版本。快速确认是否用上 GPU 的方法是import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.version.cuda) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU)如果torch.cuda.is_available()返回False最常见原因是 PyTorch 安装成了 CPU 版。需要从对应版本的源安装 CUDA 版比如 pip 指定--index-url为对应 CUDA 的下载源。Windows 下还要确认显卡驱动版本是不是够新太老的驱动会不被 CUDA 运行库支持。版本之间的差异也会影响代码。某些 API 在旧版本上不存在比如torch.amp.autocast在旧版本里叫torch.cuda.amp.autocast。写库给别人用时最好写兼容判断或者直接提示对方升级到指定的 PyTorch 版本。6.2 调试时的三个高效工具第一个是torchinfo旧名torchsummary。它能打印出网络每一层的名称、输出形状和参数量比手动一行行 print shape 高效得多。装上后一句summary(model, input_size(3, 224, 224))就能看到整个模型的结构。第二个是torch.utils.tensorboard配合可视化工具能查看 loss 曲线、梯度直方图、权重分布。我习惯在训练代码里每 100 步写一次add_scalar这样能非常直观地发现过拟合和梯度消失。第三个是torch.jit.script或onnx.export它们主要用于导出模型做推理部署。调试时可以先在 Python 里跑通再尝试导出如果导出报错说明网络中存在某些动态控制流或无法追踪的操作这能反向帮你发现 forward 里的不规范写法。我自己排查训练发散问题时有一套固定流程先看 loss 是否 nan如果是检查学习率是否过大、输入是否归一化、损失函数是否和最后的输出层匹配再看梯度范数用torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)做梯度裁剪能避免梯度爆炸最后看训练曲线是否在某个 step 突然掉下去这大概率是数据里出现了异常样本而不是模型问题。PyTorch 这个框架最大的优势是灵活但灵活也意味着你需要自己维护好这些细节。口袋参考系列第一篇先帮你把张量、自动求导和模型搭建这几块地基打牢。我写代码时始终遵循一个原则先让代码跑通再考虑优化。很多人一开始就在想怎么把代码写得很优雅结果被各种维度错误和设备错误折腾得失去耐心。你只需要把一个最简单的训练循环跑起来哪怕准确率很低也比一个永远改不完的高级框架要好得多。下次这个系列我会继续写模型保存和加载、分布式训练以及部署相关的细节。如果你在照着这部分代码实现时遇到了什么奇怪的报错建议先从 dtype 和 device 开始查因为百分之八十的问题都出在这两个地方。
阅读完成 · 觉得有帮助？
