总有人问我学 PyTorch 第一关难在哪。我的答案不是张量、不是模型结构也不是那堆层 API而是把环境装好。真的我见过太多人在配置环境这一步折腾到怀疑人生最后连import torch都没跑过就直接放弃了。这篇就聊聊我自己的实践路径不讲大而全的框架介绍只讲一件事怎么少踩坑地把 PyTorch 环境搞定以及跑通第一个训练脚本时真正需要懂的那几个核心概念。内容面向刚入门、想快速跑起来的人也适合那些已经能跑通 Demo、但对一些底层机制还一知半解的读者。1. 装环境这件事九成新手都卡在第一步1.1 先确认硬件再谈安装方式拿到一台机器我建议做的第一件事就是看显卡型号而不是先上网搜教程。PyTorch 的安装方式完全取决于你有没有 NVIDIA 显卡、驱动支持什么版本的 CUDA这步搞错了后面全是白折腾。Windows 下打开设备管理器看“显示适配器”Linux 下跑一句nvidia-smi如果这命令不存在说明驱动都没装好。这里有个常见的误区nvidia-smi显示的是驱动支持的CUDA 版本上限并不是你系统里已经装好的 CUDA 版本。很多教程会让人先去装 CUDA Toolkit但用 PyTorch 官方预编译包的话这一步基本可以跳过。你只需要一个足够新的驱动剩下的 CUDA 运行库都跟着 PyTorch 的 wheel 包一起走。如果机器没有独立显卡或者买的是某果芯片的机器那就只能走 CPU 版本。CPU 版对学习和调通代码完全够用但训练大规模模型就别想了这点需要有预期。1.2 CUDA 版本到底怎么选很多人在这卡住是因为官方安装命令里给了 cuda 12.4、cuda 12.1、cuda 11.8 好几个选项不知道选哪个。我的建议很直接看nvidia-smi顶部显示的 Driver 版本对应的 CUDA 支持范围然后选一个不低于 PyTorch 要求的版本号即可。举个例子某种驱动版本支持的 CUDA 是 12.x那你就直接选官方给的 cuda 12.4 或更高版本对应的安装命令。你不需要担心选错了会怎样只要驱动版本不落后太多PyTorch 的依赖库都能正常运行。实际上CUDA 自带向下兼容能力只要驱动支持的 CUDA 版本高于 PyTorch 所需的副本次数就没事。1.3 用 conda 还是 pip我的建议很折腾说句实话两个都能用但如果你是在自己机器上做实验我更倾向用 conda 建干净的环境。原因是 conda 不仅管 Python 包还能管一部分系统依赖库遇到坑时回滚方便。使用 pip 的话依赖关系经常把系统搞乱。conda create -n torch_env python3.10 conda activate torch_env pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124这里有个细节即便在 conda 环境里我也建议最后用 pip 装 PyTorch。因为 PyTorch 官方对 conda 频道和 pip 两套渠道的同步速度不一样pip 装的版往往比 conda 频道更新更快而且用 pip 装预编译包还省事。命令末尾那个--index-url参数就是指定从官方库的 CUDA 版本渠道下载比先装 CUDA 再装 torch 的方式要简单得多。装完第一件事跑一段验证代码import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))这三行能正确输出说明环境已经通了。记住torch.cuda.is_available()返回True才是重点版本号乱不乱无所谓。提示不要强迫症一样非要装最新版。PyTorch 生态里有些经典代码库可能还没跟上新版 API选择一个稳定版本比如 2.1 或 2.2往往更省心。2. Tensor 和 autograd真正需要搞懂的两个地基2.1 Tensor 并不神秘它就是带“加速器”的多维数组Tensor 这个词听起来高大上但你完全可以把它理解为带各种附属信息的多维数组。它和 NumPy 的 ndarray 很像区别在于 Tensor 能自动在 GPU 上做运算还能记录计算过程。举个我常给朋友讲的例子把 Tensor 想象成一张快递单。它本身装着数据好比包裹还附着一些“物流信息”——在哪个设备上CPU 还是 GPU、需要不需要记录梯度requires_grad、梯度的值是什么grad。普通数组就像你自己拎着包裹走Tensor 则相当于走了一套完整的物流系统每一步都有迹可循。在代码层面最常用的创建方式有几种import torch # 从列表创建 x torch.tensor([1.0, 2.0, 3.0]) # 全零 / 全一 / 随机 y torch.zeros(2, 3) z torch.ones(2, 3) w torch.randn(2, 3) # 像 NumPy 一样做运算 result x.sum()Tensor 默认在 CPU 上想挪到 GPU 就调用.to(cuda)。不过别急着把每个数据都手动搬来搬去后面用DataLoader和.to(device)的模式会更规范。2.2 autograd 是怎么做到“不用手算梯度”的这是初学 PyTorch 时最值得多花时间理解的一块。传统方法里你要自己推导出每一层参数的梯度公式不仅繁琐还容易出错。PyTorch 的自动求导机制干的事是把你的计算过程记录成一张图然后按链式法则反过来算梯度。可以把这个过程想象成顺着水管走一遍再逆着水管走一遍。前向传播的时候水从进口流向出口反向传播的时候水流的方向反过来每个阀门参数都能知道自己被推了多少力。这就是为什么只要设置requires_gradTrue然后在 loss 上调用.backward()所有参数的.grad就自动填上了值。有个点很多人会忽略默认情况下Tensor 的requires_grad是False。只有当你打算对这个张量求梯度才需要设成True。对模型参数来说nn.Parameter会默认开启这个属性所以平时你不太需要手动管。再看一段简单的代码x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # 结果应该是 7因为 2x3 在 x2 时是 7这代码虽然简单但背后的意义不简单——对学习率、反向传播流程的理解全靠这个机制打底。2.3 用 autograd 跑通第一个最小示例不急着上模型库先用最原始的方式手写一层逻辑这样理解最深。假设我们要拟合一条直线y w * x bimport torch w torch.randn(1, requires_gradTrue) b torch.zeros(1, requires_gradTrue) x_data torch.tensor([1.0, 2.0, 3.0, 4.0]) y_data torch.tensor([2.0, 4.0, 6.0, 8.0]) learning_rate 0.01 for epoch in range(100): y_pred w * x_data b loss ((y_pred - y_data) ** 2).mean() loss.backward() with torch.no_grad(): w - learning_rate * w.grad b - learning_rate * b.grad w.grad.zero_() b.grad.zero_() if epoch % 20 0: print(fepoch {epoch}: loss{loss.item():.4f})这个代码里有两个值得留意的操作一个是torch.no_grad()它告诉 PyTorch 这段操作不需要记录计算图因为参数的更新本身不该被纳入梯度计算另一个是每次更新后要调zero_()清空梯度不然下一次backward()会把新旧梯度累加起来结果就错了。把这段跑通你会对训练的本质有体感计算预测值、算 loss、反向传播、更新参数、清零梯度这个过程循环往复。之后再用nn.Linear这些高层 API也只是把这个流程封装了而已。3. 从玩具代码到能跑的脚本数据管道和训练流程的细节3.1 DataLoader别小看数据入口的写法很多新手一开始图省事把整批数据一次性加载进内存然后手动切 batch。数据量小的时候没什么感觉数据一多就直接内存爆掉。PyTorch 提供了一套标准的数据处理流程核心是Dataset和DataLoader两个类。你只需要定义好一个继承自Dataset的类实现__len__和__getitem__两个方法剩下的 batch 分割、乱序、多进程读取都由DataLoader处理。这样做的好处很明显数据量大时不用一次全读入内存按需读取即可。import torch from torch.utils.data import Dataset, DataLoader import torchvision.transforms as T class MyDataset(Dataset): def __init__(self, data, targets): self.data data self.targets targets def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.targets[idx] dataset MyDataset(...) dataloader DataLoader( dataset, batch_size32, shuffleTrue, num_workers2, )batch_size是每次迭代喂给模型多少条样本shuffle是在每个 epoch 开始时打乱数据顺序num_workers是开启几个子进程来加载数据。需要注意的是在 Windows 上num_workers设置过高有时反而会拖慢速度还容易触发奇怪的报错建议先设成 0 或 1 验证逻辑再逐步往上调。3.2 训练循环的工程细节从写研究 Demo 到写工程代码最大的变化是不要把关键超参数写死在代码里。学习率、batch size、epoch 数、模型保存路径这些东西最好集中放在一个配置里甚至可以直接用命令行参数控制。这一招在反复调参、复现实验时特别管用。一个标准的训练循环长这样model SimpleNet() optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion torch.nn.CrossEntropyLoss() for epoch in range(num_epochs): model.train() total_loss 0.0 for batch_idx, (data, targets) in enumerate(train_loader): data, targets data.to(device), targets.to(device) optimizer.zero_grad() outputs model(data) loss criterion(outputs, targets) loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}/{num_epochs}, Loss: {total_loss / len(train_loader):.4f})这里的model.train()不是摆设它会把模型切换到训练模式。部分层比如 Dropout 和 BatchNorm在训练和推理时的行为不一样你不调用这个实验结果往往会莫名奇妙地变差。同理在推理阶段要调model.eval()并且建议用torch.no_grad()包裹推理过程省显存还加速。3.3 验证集不是可选项很多人图省事直接拿着训练集上的 loss 判断模型好坏。这是相当危险的做法。训练 loss 低只能说明模型记住了这批数据并不能说明它对没见过的数据表现如何。正确做法是从训练集里划出一部分做验证集每个 epoch 结束后在验证集上评估一次用它作为模型好坏的标尺。划分方法很简单比如用torch.utils.data.random_split从数据集里随机取 80% 做训练20% 做验证。验证集的评估代码基本就是推理代码注意要加model.eval()和torch.no_grad()。验证 loss 的价值在于它能暴露“过拟合”的苗头。当训练 loss 还在往下降、验证 loss 开始回升的时候就是模型在死记硬背数据而不是学习规律了。见到这种信号果断停掉训练不要跟它死磕。4. 显存溢出和 loss 不降两个让人抓狂的实战问题4.1 OOM 除了调小 batch 还能怎么办显卡显存溢出是几乎每个人都会碰到的报错提示信息很明确一句 CUDA out of memory 就给你计划好的下午茶时间取消了。最直接的解决办法是把batch_size调小但这不是唯一的思路。我的处理顺序通常是这么来的关掉其它占显存的应用比如浏览器里大量标签页。检查代码里是不是有不该留在显存的张量用了del之后配合torch.cuda.empty_cache()清一下缓存。调小batch_size操作系统会自动用更小的内存占用计算梯度。如果是训练大模型可以把optimizer.zero_grad(set_to_noneTrue)配合loss.backward()中间的中间结果释放。如果模型本身太大可以考虑用混合精度训练。PyTorch 官方提供的自动混合精度接口能把刻意不需要高精度的部分用半精度计算显存占用能明显下降而且速度往往更快。需要注意的是torch.cuda.empty_cache()只是清空缓存并不是让占用完全归零。真正的问题还是要在代码层面减少驻留在显存里的张量数量。4.2 loss 不降甚至变成 NaN遇到 loss 不降不要急着调参先判断是哪种不降第一类loss 一直非常大而且不稳定甚至出现 NaN。这种情况通常是梯度爆炸了。常见诱因是学习率太大、模型初始化不好、或者是 loss 计算里某个值变成无穷大。检查方法是在训练循环里打印梯度范数total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.detach().data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5如果梯度范数在某个 step 突然飙到非常大基本就锁定了罪魁祸首。解决办法是降低学习率或者用梯度裁剪把这些梯度值限制在一定范围内。第二类loss 一直平稳降得非常慢或者干脆不动。这有可能是学习率太小也有可能是模型太简单、表达能力不足。可以尝试把学习率提高一个数量级看反应。如果提高了还是没动静问题可能出在数据预处理上比如没有做归一化输入特征数值范围差距太大导致优化路径极其曲折。第三类loss 正常下降但在某个水平就卡住。这种情况往往是模型容量到了瓶颈或者遇到局部最优。可以试试调整层数、宽度或者改用更好的优化器比如把朴素 SGD 换成 Adam 或 AdamW。我自己的经验法则是先排除数据问题再检查计算图逻辑最后才动优化器参数。大部分“玄学”的 loss 问题最后都出在这些基础环节上。4.3 保存报错时的检查清单排在训练中经常遇到的问题里还有不少是“从入门到放弃”级别的报错。比如维度不匹配的错误往往是__getitem__返回的数据形状和模型输入的尺寸对不上。出这种错别硬猜把每次 batch 的 shape 先在DataLoader里打印出来看一眼。再比如设备不一致的错误模型在 GPU 上数据还在 CPU 上一算就报错。我的习惯是在数据 batch 出来后直接.to(device)并且保证模型也在同一 device 上不要一个地方写死cuda另一个地方写cpu。还有一个容易被忽视的问题跨 epoch 的随机种子。如果你的代码里用了随机性操作又没有固定随机种子那么每次实验结果可能差异都很大。做实验可以但做复现就麻烦了这一点接下来展开聊。5. 可复现性和断点续训让实验最让人省心的习惯5.1 固定随机种子深度学习过程里有三处随机数据加载打乱的随机、模型参数初始化的随机、以及某些操作本身比如 Dropout的随机。为了让实验可以对着比有必要在训练前固定随机种子import random import numpy as np import torch seed 42 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed)如果有 GPU还想尽量复现可以再加一句torch.cuda.manual_seed_all(seed)需要说明的是不管怎么固定GPU 上某些操作的运算顺序仍然可能存在微小的不确定性导致结果不能 100% 一致。但对于大多数实验来说固定种子之后的结果已经具有可比性了不要为了追求绝对复现把自己搞疯。如果你的数据加载用了DataLoader还要确保shuffle的随机状态在每次 epoch 也是可控的通常设置DataLoader的generator参数就能做到。不过我个人的经验是只要固定了总种子大部分情况下已经够用了除非你要做严格对比实验。5.2 简单的 checkpoint 保存与恢复训练 10 个小时结果中途断电或者显存炸一下全部白费这种事我经历过太多次后来一定会写 checkpoint 逻辑。核心就是把模型状态、优化器状态、当前 epoch、当时的 best loss 都存到一个文件里checkpoint { model_state: model.state_dict(), optimizer_state: optimizer.state_dict(), epoch: epoch, best_loss: best_loss, } torch.save(checkpoint, checkpoint.pt)恢复的时候这样写checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state]) optimizer.load_state_dict(checkpoint[optimizer_state]) start_epoch checkpoint[epoch] 1 best_loss checkpoint[best_loss]这里有几个关键点优化器状态必须一起保存否则恢复训练时优化器的动量信息会丢失学习率调度也会对不上相当于从零开始如果你用了学习率调度器比如StepLR最好也把它的状态一起存进去。判断一个 checkpoint 写得好不好就看一件事——断电之后能恢复到几乎一模一样的训练状态。我习惯的做法是每训练完一个 epoch如果验证 loss 比之前最好记录更低就覆盖保存一次。这样磁盘上永远只保留最好的模型而不是堆一堆没用的中间文件。5.3 记录实验数据的习惯最后说个容易被忽略但又很重要的习惯把每次实验的关键信息记录下来。我用一个简单的表格记录内容包括数据集取哪一份、模型结构、学习率、batch size、seed、最终验证精度、跑了多少 epoch、还有什么异常现象。看上去有点原始但真到了你要回头找“上次那个效果好的实验到底用的什么配置”的时候你会发现比任何花哨的工具都可靠。PyTorch 生态里有很多现成的实验管理工具但我始终觉得先把本地手动记录这件事做好比工具本身更重要。工具可以提高效率却替代不了对实验过程的把握。我自己被“这个效果好像复制不出来”坑过很多次最后发现都是因为当时没把 seed 和版本号记下来。现在哪怕只是临时做一个快速实验文件夹里也会随手放一个config.json把该记录的参数都写上。这个小习惯比任何优化技巧都更能省时间。