★ FEATURED ARTICLE 1. 为什么我最终选择了 PyTorch聊 PyTorch 之前先说说我自己的背景最早做深度学习用的是某静态图框架写一个模型要先定义好整个计算图再喂数据进去跑。那时候最痛苦的不是写网络本身而是调试。一旦计算图构建阶段报错密密麻麻的报错信息指向的是编译层的节点你根本不知道是自己的维度写错了还是逻辑有问题。后来接触到 PyTorch第一次感受到什么叫做“写代码像写普通 Python 一样”。模型里每个操作都是即时执行print 能直接打印张量的形状和值断点打在 forward 里看到的变量就是真实的中间结果。这种体验上的差异直接改变了我的开发习惯。PyTorch 到底是什么它本质上是一个基于 Python 的科学计算库核心数据结构是张量Tensor配合自动求导机制Autograd可以高效地完成神经网络的定义、训练和部署。它的动态计算图特性让模型的修改、实验的迭代变得极其灵活——你在 forward 里写 if 语句、写 for 循环、用 Python 的任意控制流都不会有任何限制因为每次前向传播都是重新构建计算图。这就好比你在做菜时随时调整调料用量而不是提前把所有步骤都写死在流程表里。有人会问现在框架这么多为什么非要选 PyTorch我个人的理由有三个第一学术界的默认选择几乎每天都能看到基于 PyTorch 的新论文和开源代码你想复现别人的工作直接 pip install 就能跑不用把模型从另一个框架翻过来第二调试体验无敌torch 张量和 NumPy 数组可以无缝互转配合 Python 生态里的各种可视化工具做分析非常顺手第三从研究到生产的路径短TorchScript、ONNX 导出、TorchServe 这些工具链让模型部署变得比想象中顺畅。这篇文章不是 PyTorch 的官方文档翻译而是我实际使用几年后总结的实战笔记。我会从框架选型的思考、核心概念拆解、完整训练流程搭建、到踩坑实录一层层展开。不管你是有一些深度学习基础但还没上过手还是已经在用 PyTorch 但想系统理清底层机制都能在这篇里找到对应的内容。2. 核心概念拆解张量、自动求导与模块化2.1 张量不仅仅是多维数组PyTorch 里的一切操作都是围绕张量Tensor展开的。你可以把张量理解成 NumPy 的 ndarray 加上两个附加能力GPU 加速计算和自动求导跟踪。初学者最容易忽略的一点是torch.Tensor 和 numpy.ndarray 虽然长得像但内存布局和底层实现并不相同。很多坑正是从“我直接把 numpy 的代码搬过来”开始的。创建一个张量很简单常用的方式有这么几种import torch # 从列表创建 t1 torch.tensor([1.0, 2.0, 3.0]) # 全零张量 t2 torch.zeros(3, 4) # 随机张量 t3 torch.randn(2, 3) # 和 numpy 互转 import numpy as np arr np.array([[1, 2], [3, 4]]) t4 torch.from_numpy(arr) # 注意共享内存 arr2 t4.numpy() # 反向转换同样共享内存这里有个很重要的细节torch.from_numpy()转换出来的张量与原来的 NumPy 数组共享底层内存。你改了数组的值张量的值也会跟着变。这在做数据预处理时很方便但如果你没意识到这一点在训练循环里无意修改了原始数组会引入特别隐蔽的 bug。我踩过一次数据归一化时直接对原始数组做 in-place 操作结果训练过程中验证集的指标每隔几个 epoch 就跳一下后来排查了半天才发现是共享内存导致的。张量的另一个核心操作维度是设备device。CPU 和 GPU 上的张量是两个世界的生物相互之间不能直接做运算。你写a b时如果 a 在 GPU、b 在 CPU会直接报错不会自动帮你迁移。这个设计看起来不友好实际上是为了让你明确知道数据在哪避免在训练循环里因为隐式的 CPU-GPU 拷贝导致性能骤降。2.2 自动求导grad_fn 与反向传播的机制PyTorch 最核心的机制就是自动求导。你不需要手动推导链式法则也不需要自己写反向传播代码只需要在创建张量时指定requires_gradTruePyTorch 就会在每次前向传播时记录所有操作形成一个计算图然后在你调用backward()时自动计算梯度并填充到每个叶子节点的.grad属性中。看一段最小示例x torch.tensor(2.0, requires_gradTrue) y x ** 2 3 * x y.backward() print(x.grad) # tensor(7.)因为 dy/dx 2x 3x 2 时导数为 7结果符合预期。这个过程背后有两个关键属性grad_fn表示该张量是由什么操作产生的is_leaf表示该张量是否为叶子节点。只有叶子节点通常是你直接创建的、requires_gradTrue 的输入才会有梯度累积。实际训练中这一机制最需要注意的地方有三个。第一个是梯度累积默认情况下调用backward()后梯度是累加的而不是清零的。这意味着如果你在同一个 batch 上连续调用两次backward()梯度会叠加在一起。所以每个 batch 开始前必须调用optimizer.zero_grad()或者等价地model.zero_grad()否则梯度会一直累加下去训练会出现各种诡异的不收敛现象。第二个是torch.no_grad()在推理阶段、计算验证集指标时我们不需要梯度此时应该用with torch.no_grad():包裹代码块。这能显著减少显存占用和计算开销因为 PyTorch 不再需要记录任何操作来构建计算图。如果不加我实测一个小模型在验证阶段显存开销能差出 30% 以上。第三个是detach()方法当你想从一个计算图中取出某个张量用于其他计算但又不想让梯度传播到它来源的路径上时调用tensor.detach()会返回一个新的张量它和原张量共享数据但不参与梯度追踪。这在做类似 GAN 的训练、或者对中间特征做可视化时非常常用。忘了 detach 导致梯度流错方向其实是很多高级训练程序里出现 NaN 或数值爆炸的根源。2.3 nn.Module把模型的每个零件组装起来torch.nn.Module是构建神经网络的基础类。你定义任何一个模型都会继承这个类然后在__init__里声明层和参数在forward里定义前向计算逻辑。为什么 PyTorch 要这么设计因为nn.Module提供了几个极为重要的隐式功能自动追踪参数你在__init__中赋值的任何nn.Parameter、任何子模块都会被自动收集到model.parameters()里。接下来优化器、存储、设备迁移等操作都不需要你手动指定每层参数。model.to(device)一键迁移调用model.cuda()或model.to(cuda)时模型所有参数和缓冲区都会被递归地迁移到目标设备。model.train()/model.eval()模式切换这会影响到 BatchNorm、Dropout 等层的行为。训练时要开 Dropout验证时要关掉这个切换是全局的。一个简单的两层模型定义如下import torch.nn as nn class TwoLayerNet(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 TwoLayerNet(784, 256, 10) print(model)这里有个新手很容易踩的误区在__init__里用普通 Python 列表list保存子模块是无效的。nn.ModuleList和nn.Sequential才是正确的容器因为它们会让 PyTorch 正确地注册模块和参数。我自己早期写过用 list 保存三层全连接结果model.parameters()一个参数都拿不到优化器傻眼了梯度也算不上来训练直接卡死。nn.Sequential则是快速堆叠层的方式适合没有复杂分支的直筒型网络model nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) )相比之下如果你需要实现类似 ResNet 那样的残差连接、或者 Multi-Branch 结构就必须用自定义nn.Module的写法。forward里写什么都是合法的——包括条件分支、循环、甚至调用外部函数这完全体现了动态图框架的灵活性。3. 从零搭建一个完整训练流程3.1 数据加载Dataset 与 DataLoader 的使用逻辑训练任何模型第一步都是把数据组织好。PyTorch 的torch.utils.data.Dataset是一个抽象类你只需要实现两个方法__len__返回样本总数__getitem__根据索引返回一个样本。这样设计的好处在于数据读取可以做到懒加载——不会一次性把所有数据读进内存而是在迭代时按需读取。写一个最小示例假设我们要处理图片分类任务图片存放在文件夹里标签在 CSV 文件中from torch.utils.data import Dataset, DataLoader import pandas as pd from PIL import Image class ImageDataset(Dataset): def __init__(self, csv_path, img_dir, transformNone): self.df pd.read_csv(csv_path) self.img_dir img_dir self.transform transform def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img_path f{self.img_dir}/{row[filename]} image Image.open(img_path).convert(RGB) label row[label] if self.transform: image self.transform(image) return image, label dataset ImageDataset(train.csv, images/) dataloader DataLoader(dataset, batch_size32, shuffleTrue, num_workers4)DataLoader 是真正在训练循环里被迭代的东西。它把 Dataset 里的样本按 batch 大小打包并支持多进程并行读取、打乱顺序、pin_memory 加速等。几个关键参数的含义如下batch_size每个 batch 的样本数直接影响显存占用和梯度更新频率。shuffle每个 epoch 是否打乱样本顺序。训练集一般设为 True验证集设为 False。num_workers开启几个子进程来加载数据。增大这个值可以显著缓解 CPU 数据读取和 GPU 计算速度不匹配的问题但不是越大越好过大的进程切换开销反而会拖慢速度。pin_memory设为 True 时如果你的机器使用了 CUDADataLoader 会将数据放在锁页内存中这个细节能提升 CPU 到 GPU 的拷贝速度。这里有一个经常被忽略的坑Dataset 的__getitem__里如果做了复杂的图像解码和预处理那么每个进程都会有一份 Dataset 的副本占用额外内存。如果你的数据量特别大或者预处理特别重建议把耗时部分在__getitem__里做缓存或者用内存映射的方式避免每个 worker 都重复加载同一份数据。3.2 模型定义与参数初始化数据准备好了接下来是模型。模型定义方法在上面已经讲过这里重点说说参数初始化的细节——这是很多人忽略但影响很大的环节。PyTorch 的nn.Linear、nn.Conv2d等层在创建时有默认的初始化策略如 kaiming_uniform。在大多数情况下默认初始化已经够用你不需要额外做什么。但是当你使用一些自定义层、或者参考某些论文的实现需要特定初始化策略时可以通过nn.init来操作def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_normal_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) model.apply(init_weights)model.apply()会递归地遍历所有子模块对每个模块调用传入的函数这是统一的初始化入口。为什么不建议在__init__里逐层手动初始化因为当你的模型有几十层时手动写会非常冗长而且容易漏掉某层。初始化策略怎么选简单说几个常见的xavierGlorot系列适合使用 Sigmoid / Tanh 等饱和激活函数的网络。kaimingHe系列适合使用 ReLU 及其变体激活函数的网络也是 ResNet、VGG 等经典结构常用的策略。全零初始化一般不要用会导致所有神经元对称更新方向一致无法学到差异化特征。3.3 损失函数与优化器的选择PyTorch 在torch.nn里提供了大量损失函数torch.optim里提供了常见优化器。这里不打算罗列所有 API而是聊一聊选型时的实际考虑。对于分类任务最常用的是nn.CrossEntropyLoss()。注意这个函数默认接受的是原始 logits即网络最后一层没有过 softmax 的输出它在内部会自行计算 softmax。所以你不需要在网络的forward里额外加nn.Softmax()否则会做一次没有意义的 softmax 再 softmax既浪费计算量还可能影响数值稳定性。对于回归任务通常用nn.MSELoss()。如果你发现回归训练输出的数值始终不稳定检查一下数据是否有做标准化、预测值是否在合理范围内。损失值一上来就 NaN 的常见原因之一就是目标值尺度过大。优化器方面工程上我从 SGD带动量和 Adam 聊起。“带动量 SGD”是最经典的选择收敛稳定、可解释性强很多传统 CNN 任务用它能得到比 Adam 更好的泛化效果Adam 是自适应学习率方法调参省心对超参数不那么敏感适合做 baseline 和新任务的第一版实验。现在还有不少人在用 AdamW解耦权重衰减它在 Transformer 系列模型中是标配。我的实操建议是新任务先上 AdamW学习率用默认的 1e-3 左右跑一个 baseline确认整个流程能跑通。然后如果想追求更好的效果再切到 SGD Momentum数值一般设 0.9配合学习率衰减策略StepLR / CosineAnnealingLR精调。这样既不浪费时间又能有一个较优的结果托底。3.4 完整的训练循环示例把前面几个模块组装起来完整的训练代码骨架如下。我用一个在 MNIST 上的分类任务作为例子但结构完全适用于其他任务import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 数据预处理和加载 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse, num_workers2) # 模型、损失函数、优化器 model nn.Sequential( nn.Flatten(), nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 10) ) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-3) # 训练 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) def train_one_epoch(loader, model, criterion, optimizer, device): model.train() total_loss 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) return total_loss / len(loader.dataset) def evaluate(loader, model, device): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return correct / total for epoch in range(10): train_loss train_one_epoch(train_loader, model, criterion, optimizer, device) test_acc evaluate(test_loader, model, device) print(fEpoch {epoch1}: loss{train_loss:.4f}, acc{test_acc:.4f})这个代码骨架有几个细节值得解释。model.train()和model.eval()是成对出现的。训练模式下 Dropout 启用、BatchNorm 更新 running mean/varianceeval 模式下这些行为被关闭模型用累积的统计量做推理。漏掉train/eval切换大概是新手中最常见的问题之一——训练指标正常验证指标却很奇怪很可能就是这里出了问题。optimizer.zero_grad()必须放在每次loss.backward()之前。我习惯把它放在前向传播之前调用即每个 batch 开始时就清理梯度。有些人写在loss.backward()之后逻辑上没问题但一旦代码里有断点或者分支提前 continue 了反而会漏清。with torch.no_grad()是验证阶段的标准写法。它会阻止 PyTorch 为验证阶段的操作构建计算图省显存、省时间。等你跑的数据量大、模型又深的时候这个细节带来的差别非常明显。outputs.argmax(dim1)是获取预测类别的常见写法。CrossEntropyLoss 的输出维度是 (batch, num_classes)在每个样本维度上取最大值对应的索引就是模型的预测类别。4. 实战中常见的坑与排查技巧实录4.1 显存相关OutOfMemory 的几个隐藏元凶CUDA OutOfMemory 是 PyTorch 用户最常碰到的错误之一。最常见的原因当然是 batch size 太大、模型太大、或者分辨率太高。但有些更隐蔽的情况我实际排查时发现不只一个人在这上面卡过。第一个隐藏场景是验证阶段没有开torch.no_grad()。如果你在验证循环里忘了包这层每个验证 batch 都会临时构建一个计算图而这些计算图只有在变量被销毁时才会释放。如果验证集很大或者你验证时也跑了很多次迭代显存会被一点一点吃光。我的经验是训练循环已经把显存占得很满了验证阶段哪怕只多几个 batch 的计算图也可能直接触发 OOM。第二个场景是梯度累积过程中产生了不必要的中间变量。在训练循环里如果你写了类似loss criterion(outputs, labels) some_regularization(outputs)这样的代码而some_regularization的结果本身不需要梯度就最好detach掉。否则 PyTorch 会把整条计算路径的记录都留着白白消耗显存。第三个场景是 DataLoader 的 worker 和主进程之间的数据缓冲。num_workers设得过高每个 worker 的预取缓冲会在 CPU 内存上吃很多资源虽然不直接占用 GPU 显存但可能会让系统整体变慢间接影响 GPU 显存的释放节奏。如果你发现 OOM 只出现在训练中途而不是一开始检查一下是不是系统内存不够用导致的数据交换问题。第四个非常隐蔽的坑是使用torch.Tensor还是nn.Parameter的混淆。在自定义层里如果你把一个非参数张量注册成 buffer比如 BatchNorm 的 running_mean默认情况下它是不需要梯度的但会占用显存。如果你在训练循环里不小心对 buffer 做了需要梯度的操作计算图也会被记录下来。4.2 设备不一致导致的“隐形崩溃”RuntimeError: Expected all tensors to be on the same device是另一个高频报错。报错信息本身其实已经说明白了就是某两个参与运算的张量一个在 CUDA 一个在 CPU。但难办的是当你的代码里有几十个张量时到底哪个在 CPU 上不是一眼能看出来的。我的排查步骤是先在报错的地方打断点逐个print(tensor.device)看看哪些张量在 CPU然后往上游追看这个张量是在哪里产生的。最常见的源头有三个数据加载后忘了images.to(device)。模型本身的参数没有调用model.to(device)。一个张量来源于 NumPy 数组或 Python list通过torch.tensor()创建后直接参与计算默认就在 CPU 上。还有一个容易被忽视的细节两个张量一个在cuda:0一个在cuda:1虽然都是 GPU但因为不在同一张卡上同样会报错。做多卡训练时要确保所有参与运算的张量都在同一个设备上。简单粗暴的做法是代码里统一用一个device变量所有.to(device)都传同一个值避免用手写的cuda字符串分散在项目各处。养成一个习惯模型定义完立刻model.to(device)数据进入训练循环第一行就images, labels images.to(device), labels.to(device)这样能避免 90% 的设备问题。4.3 梯度相关NaN、爆炸和错误累积训练过程中 loss 变成 NaN 或者 inf是另一个让人头痛的问题。Loss 是 NaN 的常见原因我在实际项目里遇到过这几种第一种梯度爆炸。特别常见于 Transformer 结构的训练深层网络的梯度相乘会导致数值溢出。解决方案是梯度裁剪gradient clipping。PyTorch 里就一行torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)放在 optimizer.step() 之前调用。它会将所有参数的梯度整体缩放使得梯度的总范数不超过阈值从而防止极端梯度值导致的参数更新异常。第二种学习率过大。模型开始几步更新后参数直接飞出去loss 瞬间变成 NaN。这种情况可以将学习率调小一个数量级试试。用 Adam 的话1e-4 到 3e-4 是很多视觉任务的安全区间。第三种损失函数或数据里有 NaN。举个具体例子如果输入数据包含 NaN比如 CSV 文件读到空值填充为 NaN前向传播后 loss 自然就是 NaN梯度自然也是 NaN。网上有一些项目会默认数据是干净的但现实中脏数据才是常态。在数据加载环节加入torch.isnan(tensor).any()检查很有效if torch.isnan(images).any() or torch.isinf(images).any(): print(f发现异常数据样本索引 {idx})第四种混合精度训练时的数值溢出。FP16 的取值范围比 FP32 小得多累加梯度时容易溢出。PyTorch 的torch.cuda.amp.GradScaler就是为了解决这个问题存在的。如果你自己手动把模型转成半精度而没有用 AMP 的梯度缩放机制遇到 NaN 是早晚的事。梯度累积还有一种情况是“看似正常的错误”梯度累积方向不对。举个例子某个参数在反向传播后不应该有梯度但由于计算图里有一条不该有的路径连着它导致它收到了梯度训练行为就会变得不可解释。排查这种问题要靠requires_grad和各张量的grad_fn来检查。如果发现某个中间张量意外带上了梯度依赖检查它是不是应该detach()。4.4 常见问题速查表现象最可能的原因快速解决方法训练 loss 正常验证 loss 反而一直下降不了忘了model.eval()切换验证循环前加上model.eval()完事后再model.train()每个 epoch 开始几个 batch 很慢后面变快DataLoader num_workers 太小数据读取瓶颈增大 num_workers开启 pin_memory验证时显存爆炸没有用torch.no_grad()验证循环用with torch.no_grad():包裹模型参数没更新loss 纹丝不动optimizer.zero_grad()忘了调或者优化器和模型参数不匹配检查每个 batch 是否调用了 zero_grad检查优化器是否传入了model.parameters()在nn.Module里用 list 保存层报参数不存在list 不会被注册为子模块改用nn.ModuleList或nn.Sequential提示两个张量 device 不一致数据/模型没有.to(device)统一使用 device 变量所有张量显式迁移保存模型后加载报错键名不匹配state_dict和当前模型结构不一致加载前打印state_dict.keys()与model.state_dict().keys()对比多线程训练速度不增反降num_workers 过高导致进程切换开销过大一般设为 CPU 核心数的一半左右再实测调整4.5 模型保存与加载的两种方式模型保存看起来简单但实际中有很多人在跨设备、跨版本时踩坑。PyTorch 保存模型有两条路线整模型保存和 state_dict 保存。整模型保存torch.save(model, model.pt)这种做法直接保存整个模型对象加载时一行model torch.load(model.pt)就能用。但缺点也很明显模型结构被固化在文件里代码里如果改了模型类名或路径加载很容易出问题。而且在不同版本的 PyTorch 之间整模型保存的兼容性不佳。更推荐的做法是保存state_dicttorch.save(model.state_dict(), model_weights.pt) # 加载 model TwoLayerNet(784, 256, 10) model.load_state_dict(torch.load(model_weights.pt))这种方式只保存参数字典不保存结构加载时需要先实例化一个结构完全一致的模型再把参数灌进去。优点在于跨 PyTorch 版本兼容性更好加载逻辑更明确还能配合不同模型结构做参数迁移。跨 GPU/CPU 加载时注意map_location参数model.load_state_dict(torch.load(model_weights.pt, map_locationcpu))如果你在 GPU 上训练保存的模型要在没有 GPU 的机器上加载不指定map_location会直接报错。同理保存时如果用了 DataParallelstate_dict里的键名会多出module.前缀加载到单卡模型前要去掉这个前缀或者加载后手动替换键名state_dict torch.load(model_weights.pt) from collections import OrderedDict new_state_dict OrderedDict() for k, v in state_dict.items(): name k[7:] if k.startswith(module.) else k # 去掉 module. 前缀 new_state_dict[name] v model.load_state_dict(new_state_dict)这个细节我第一次遇到时排查了快一个小时报错信息显示键名不匹配数下来数量是一样的但对不上号。后来一查才知道是 DataParallel 加上去的前缀。5. 提升效率的几个实用技巧5.1 用 AMP 自动混合精度训练减少显存占用显存不够用最立竿见影的手段就是混合精度训练。PyTorch 的自动混合精度AMP功能可以让你在几乎不损失精度的情况下把计算量降低一个档次显存占用也明显减少。原理很简单在适当的环节使用 FP16半精度计算在关键环节保留 FP32并通过梯度缩放来避免 FP16 的数值溢出。实际使用方法非常直接from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()加上这几行以后我的一个小型模型训练时显存占用从约 8GB 明显下降到 5GB 左右训练速度也快了一些。遇到大模型或者大 batch 需求时这是最值得优先尝试的优化方案。5.2 使用 TensorBoard 记录训练过程调试模型时只打印 print 日志远远不够尤其是观察 loss 曲线和权重分布的变化时可视化工具的价值就体现出来了。PyTorch 官方提供了torch.utils.tensorboard接口不需要额外装插件直接写代码记录即可from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/experiment_1) # 每个 epoch 结束 writer.add_scalar(loss/train, train_loss, epoch) writer.add_scalar(acc/test, test_acc, epoch) # 可视化模型结构 dummy_input torch.randn(1, 28 * 28, devicedevice) writer.add_graph(model, dummy_input) writer.close()训练完后在项目目录运行tensorboard --logdirruns浏览器打开指定地址就能看到 loss 曲线和模型结构。add_scalar和add_graph是最常用的两个接口后者可以直接可视化模型计算图对于排查“为什么梯度传不到第一层”这类问题时特别有用。5.3 让训练循环更有条理本地代码组织建议PyTorch 本身对项目结构没有任何强制要求但实际工程里一个好组织方式能省下大量调试时间。我做项目时习惯把代码拆成这样├── config.py # 超参数、路径配置 ├── dataset.py # Dataset 定义 ├── model.py # 网络结构定义 ├── train.py # 训练循环 ├── utils.py # 指标计算、可视化等辅助函数 └── runs/ # TensorBoard 日志输出目录每个文件只做一件事模型结构单独放一块训练逻辑单独放一块。当模型结构需要调整时只改model.py当想换一个数据增强方案时只改dataset.py。这种解耦不是说教而是我自己经历过把模型定义和训练逻辑写在一个大文件里后期改结构时小心翼翼怕碰坏其他代码的狼狈状况。有一点要强调超参数最好集中写在config.py里不要散落在各个代码文件中。这样当你想跑一组对比实验时只需要修改配置文件中的值而不用在代码里搜索lr 出现几次。我早期吃过这个亏——同一个学习率在三个文件里各写了一次改了一个漏了另外两个结果一组实验花了一周时间跑完后才发现参数并没有全部生效。根据我个人的实际体会PyTorch 让我坚持用下来的最大理由不是某个单一功能而是它把“快速试错”这件事做到了极致。框架层面的自由和生态的富饶意味着你在碰到任何具体问题时总能找到现成的轮子或者前人的方案。上面这些内容都是我在一个个实际项目里真金白银换来的经验希望能帮你少走一些弯路。如果你刚开始学照着第三章的代码骨架跑通一两个小项目再回来体会第四章里的那些坑你会对 PyTorch 的理解上一个台阶。
阅读完成 · 觉得有帮助？
