2026/8/30 2:29:08

PyTorch入门实战:环境搭建、张量、自动求导与混合精度

PyTorch入门实战:环境搭建、张量、自动求导与混合精度 PyTorch 是深度学习研究和工程落地里绕不开的框架。无论是跑视觉模型、文本模型还是想自己验证一个新想法PyTorch 的动态图和自动求导机制都会让代码写起来更接近自然思路。这篇文章不是把官方文档重新抄一遍而是按实际使用下来最值得注意的顺序把从零上手 PyTorch 要解决的几件事讲清楚环境怎么搭、版本怎么选、核心概念怎么理解、FP32/FP16/BF16/TF32 这些精度格式在什么场景用、训练任务怎么从单条样例扩展到批量、以及报错之后先查哪里。适合刚准备入门深度学习的读者也适合已经跑过一些代码、但想把自己的训练流程整理得更规范的人。网上关于 PyTorch 的教程很多但大部分只讲了 API 长什么样没有讲清楚为什么这样写、什么情况下会出问题、低配机器怎么调整。这篇文章会把这类细节补上。学习路线不需要太长一周时间足够把一个最小的图像分类任务从环境搭建跑到模型保存前提是每一步都按顺序来不跳步。1. 先明确 PyTorch 解决什么问题一周路线怎么排1.1 PyTorch 不是“一个库”而是一整套深度学习工作流很多新手以为 PyTorch 就像 requests、numpy 一样装完 import 就能用。实际上 PyTorch 包含几层东西。第一层是张量计算库。张量可以简单理解为多维数组但和 NumPy 数组不一样的地方在于张量可以放到 GPU 上计算并且会自动记录计算过程中的操作方便后续求梯度。这是深度学习的根基没有自动求导能力反向传播算法就得手动实现。第二层是神经网络模块库。torch.nn 里封装了线性层、卷积层、循环层、损失函数、优化器等常用组件。你不需要从零写一个卷积操作直接调用nn.Conv2d就能拿到一个可训练的卷积层。第三层是数据加载工具。torch.utils.data.Dataset和DataLoader负责把原始数据组织成批次支持多进程加载、shuffle、采样器设置等。很多新手在这里栽跟头后面会单独讲。第四层是部署和导出工具。训练完模型并不意味着结束实际落地时可能需要导出成 TorchScript、ONNX 或者在终端设备上推理。这一层属于进阶内容但在一周学习结束时值得看一眼至少要知道训练和推理是两回事。理解了这四层就明白了 PyTorch 不是一个单一的库而是一整套深度学习工作流的集合。学习时不要只盯着import torch要按这个层次逐步深入。1.2 一周路线环境、单任务、批量训练、部署验证我建议把一周时间拆成四个阶段不要混着学。第一天到第二天搭环境跑通一个最小的张量例子。这一步的目标不是学会所有 API而是确认 PyTorch 能 import、能在 CPU 上计算、如果有 GPU 能识别到显卡。第三天到第四天理解自动求导和模型构建。写一个最简单的线性回归或逻辑回归把训练循环跑通。这里的关键不是模型多复杂而是要理解forward、loss、backward、optimizer.step()这些动作之间的关系。第五天到第六天做一个小型的图像分类或文本分类任务。引入 Dataset、DataLoader、验证集、模型保存和加载。这一步开始接触真实的训练流程数据怎么组织、日志怎么看、模型什么时候保存。第七天做一次部署方向的小验证。把训练好的模型保存成文件再重新加载进行推理或者导出成 ONNX 格式。同时了解一下 FP16、BF16 这些低精度格式分别在训练和推理中怎么用。这个路线看起来简单但每一步都有明显的验收标准。环境阶段的标准是torch.cuda.is_available()是否返回正确模型阶段的标准是 loss 是否在稳定下降数据阶段的标准是模型在验证集上的表现是否接近合理范围。每一步都验证通过再进入下一步。2. 环境搭建CUDA、Python、安装命令怎么选2.1 先查驱动和 CUDA 版本再装 PyTorch环境搭建是新手最容易卡住的地方但绝大多数问题其实不是 PyTorch 本身的问题而是安装前没有确认自己的硬件和系统环境。先说最普通的 CPU 环境。如果你只是想学习 PyTorch 的语法CPU 完全足够。一个简单的 MNIST 手写数字识别任务CPU 跑几十秒一个 epoch 也能接受。这种情况下只需要安装 CPU 版 PyTorch不需要关注 CUDA。如果你的电脑有 NVIDIA 显卡并且打算训练稍微大一点的模型那就要确认三样东西操作系统Windows、Linux 还是 macOS。Python 版本。PyTorch 对 Python 版本有要求一般建议使用 3.9 到 3.11 之间的版本。NVIDIA 驱动和 CUDA 版本。这里有一个很容易搞混的概念系统里的 CUDA 版本和 PyTorch 安装时选择的 CUDA 版本不一定完全一致。PyTorch 安装包本身就打包了自己的 CUDA 运行库所以即使系统只装了 NVIDIA 驱动你也可以安装带 CUDA 支持的 PyTorch。NVIDIA 驱动保持较新即可关键是安装 PyTorch 时选的 CUDA 版本要和显卡驱动兼容。在 Windows 上最简单的验证方式是打开命令行执行nvidia-smi查看右上角的 CUDA Version。这个数字表示当前驱动支持的最高 CUDA 版本。比如驱动显示支持 CUDA 12.2那么安装 PyTorch 时选择 cu121 或 cu124 这类版本一般是没问题的。2.2 安装命令里的 CPU 版和 GPU 版怎么选PyTorch 官方推荐用 pip 安装。安装 GPU 版时命令大致是pip install torch torchvision torchaudio默认情况下会安装当前平台适配的版本。如果默认源下载太慢可以换成国内镜像源或者指定 PyTorch 官方源pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121这里需要注意的是cu121表示使用 CUDA 12.1 的版本cu118表示 CUDA 11.8。选哪个不是越大越好而是要看你的驱动兼容性。如果显卡比较老驱动版本也比较老那选择 cu118 更稳妥。CPU 版安装更简单pip install torch torchvision torchaudio在 CPU 环境下默认安装的通常就是 CPU 版。如果你想要明确指定 CPU 版可以去 PyTorch 官网的安装页面选择对应配置会生成对应的完整命令。安装完成后先用一行代码验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU)如果torch.cuda.is_available()返回 True说明 GPU 环境可用。返回 False 时不要急着重装先确认安装的是 GPU 版、驱动是否正常、Python 是否为 64 位版本。2.3 conda 环境为什么要单独建我强烈建议不要直接往系统 Python 里装 PyTorch。原因是深度学习项目之间依赖很容易冲突一个项目需要 PyTorch 2.0另一个项目需要 PyTorch 1.13如果你只有一个全局环境改版本会非常痛苦。推荐用 conda 创建独立环境conda create -n pytorch_env python3.10 conda activate pytorch_env然后在这个环境里安装 PyTorch。这样不管之后是换项目还是升级版本都不会影响系统环境。创建环境时Python 版本建议选 3.10 或 3.11不要选太新的版本。PyTorch 对最新 Python 版本的适配有时会滞后选一个稳定版本能省很多事。注意如果你在安装时发现下载速度特别慢先检查是不是用的默认源。换成国内镜像源之后速度一般会明显改善。3. 核心知识点张量、自动求导、模型构建、训练循环3.1 张量从 NumPy 到 GPU 的迁移PyTorch 里最基础的数据结构是torch.Tensor。你可以把它理解成支持 GPU 加速和自动求导的 NumPy 数组。创建一个张量的方式有很多import torch # 从列表创建 t1 torch.tensor([1, 2, 3]) # 全零张量 t2 torch.zeros(2, 3) # 随机张量 t3 torch.randn(2, 3) # 和 NumPy 互转 import numpy as np arr np.array([4, 5, 6]) t4 torch.from_numpy(arr)张量和 NumPy 数组最核心的区别在于设备。你可以把张量移动到 GPU 上t_gpu t1.cuda()但前提是torch.cuda.is_available()返回 True。在 CPU 环境下调用.cuda()会直接报错所以更推荐条件判断device torch.device(cuda if torch.cuda.is_available() else cpu) t1 t1.to(device)这个device变量在后来的模型训练中会反复用到。模型参数、输入数据、标签都要放到同一个设备上否则会报设备不一致的错误。张量的另一个重要属性是requires_grad。创建张量时如果设成 TruePyTorch 会记录这个张量参与的所有计算之后可以通过反向传播计算梯度。这是理解自动求导的基础。3.2 自动求导为什么 loss.backward() 是关键传统写法里你要自己推导梯度公式然后手写更新逻辑。PyTorch 的自动求导机制把这一步自动化了你只需要把前向计算写出来PyTorch 会构建一张计算图然后在调用loss.backward()时自动计算所有参与张量的梯度。看一个最简单的例子import torch x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # 结果是 2*x 3 7这里y.backward()执行后x.grad就会被计算出来。核心规则是只有requires_gradTrue的张量才会获得梯度。在模型训练中这个机制会完整地体现出来optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step()这几行代码看起来简单但每一步都有原因。optimizer.zero_grad()必须先执行。因为梯度是累加的如果你不清空上一次的梯度下一次 backward 会把新的梯度加在旧梯度上导致更新方向错误。loss.backward()执行反向传播计算所有参数的梯度。optimizer.step()根据梯度更新参数。这几个动作的顺序不能乱。新手最容易犯的错误是忘记zero_grad()或者把step()放在backward()之前。3.3 模型构建和训练循环的最小代码构建模型有两种常见方式。一种是直接继承nn.Module另一种是使用nn.Sequential。继承nn.Module的写法更灵活适合复杂网络import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 128) self.relu nn.ReLU() self.fc2 nn.Linear(128, 10) def forward(self, x): x self.fc1(x) x self.relu(x) x self.fc2(x) return xnn.Sequential更适合简单的堆叠结构model nn.Sequential( nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10) )模型的训练循环一般长这样import torch.optim as optim model SimpleNet().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) for epoch in range(10): running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() print(fEpoch {epoch1}, Loss: {running_loss / len(train_loader):.4f})这里有几个关键判断点。第一images和labels必须和模型在同一设备上。第二loss.item()取出的是 Python 浮点数可以用来计算平均 loss。第三如果running_loss一直不下降大概率是学习率设置不对或者数据没有正确预处理。4. 浮点数格式详解FP32、FP16、BF16、TF32 怎么选4.1 深度学习训练和推理为什么需要多种精度格式很多人在跑通训练代码之后会看到一个概念混合精度训练。简单来说就是用低精度格式来加速计算、减少显存占用同时尽量保持模型效果。默认情况下PyTorch 中的模型参数和中间计算都是 FP32也就是单精度浮点数。FP32 精度高但占用的显存也大。随着模型规模增长显存成为很大的瓶颈于是大家开始使用 FP16、BF16 这些低精度格式。但低精度不是万能的。降低精度意味着数值表示范围变小、精度变低如果处理不当会导致梯度消失、loss 不收敛等问题。所以你需要理解每种格式的适用范围。4.2 四种格式的差异与适用场景FP32 是单精度浮点数用 32 位表示一个数其中 1 位符号位、8 位指数位、23 位尾数位。它在训练中精度最高但显存占用也最大。适合作为默认格式或者在没有特殊优化需求的小规模训练中使用。FP16 是半精度浮点数用 16 位表示其中 1 位符号位、5 位指数位、10 位尾数位。它的优点是显存占用减半计算速度也更快但表示范围比 FP32 小。数值比较大的梯度在 FP16 下可能溢出所以直接用 FP16 训练时常常需要配合损失缩放loss scaling。BF16 是另一种 16 位格式但它保留了 8 位指数位所以表示范围接近 FP32只是尾数位更少。它比 FP16 更适合训练因为动态范围大不容易梯度溢出。代价是精度略低。在 NVIDIA Ampere 及更新的架构上BF16 支持得很好是近期大模型训练中更常用的低精度格式。TF32 是一种特殊的格式主要用于 NVIDIA Ampere 架构上的 Tensor Core 计算。它的精度介于 FP32 和 FP16 之间计算时输入是 FP32但内部会用 19 位近似计算。TF32 不需要改代码通常在默认设置下就能生效主要适合矩阵乘法场景。4.3 实战选型建议和混合精度在实际使用中普通新手不需要一开始就手动选择精度格式。PyTorch 提供了自动混合精度工具torch.autocast和torch.amp.GradScaler。简单来说混合精度训练的做法是一部分操作用 FP16 或 BF16 计算一部分关键操作保持 FP32。这样既享受低精度的速度又尽量减少精度损失。一个典型的使用片段from torch.cuda.amp import autocast, GradScaler 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()这里注意autocast只包住了前向计算和 loss 计算反向传播由scaler.scale(loss).backward()来完成。原因是 FP16 的梯度可能过小直接反向传播会下溢所以先用 scale 放大梯度等更新参数时再缩小回来。如果你使用的是新版 PyTorch也可以更简洁from torch.amp import autocast, GradScaler格式选型时可以参考这个判断标准显存紧张且显卡支持加速优先考虑 BF16 或 FP16 混合精度。追求训练稳定性不差显存直接用 FP32。只用 GPU 推理不训练可以尝试把模型导出为 FP16 或 INT8但需要验证精度变化。显卡不支持 BF16就看硬件文档用 FP16 混合精度更稳妥。注意BF16 不是所有显卡都支持。先确认你的 GPU 架构再决定用哪种低精度格式。不要一上来就开混合精度先对比一下普通训练的 loss 曲线。5. 从 Demo 到真实任务数据处理、日志、断点保存5.1 数据加载和 DataLoader 的 batch 理解很多教程里的模型训练看起来很简单但一到真实项目就发现数据很脏、格式不统一、文件量很大。这时 Dataset 和 DataLoader 就派上用场了。Dataset 负责定义“怎么读取一条数据及其标签”DataLoader 负责把多条数据组织成批次并且支持多进程预加载。一个最简单的自定义 Dataset 长这样from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, data_list, labels): self.data_list data_list self.labels labels def __len__(self): return len(self.data_list) def __getitem__(self, idx): x self.data_list[idx] y self.labels[idx] return x, yDataLoader 使用时有两个关键参数batch_size和shuffle。loader DataLoader(dataset, batch_size32, shuffleTrue, num_workers2)batch_size表示每次模型看到多少条样本。太大会导致显存不足太小会让训练不稳定。一般从 16 或 32 开始观察显存和 loss 变化再调整。shuffleTrue会在每个 epoch 开始时打乱数据顺序避免模型学习到固定的输入顺序。验证集不需要 shuffle因为验证时不需要学习。num_workers表示用几个子进程加载数据。Windows 环境下有时候这个参数会导致问题如果报错可以先把它设成 0 验证。5.2 训练日志和模型保存训练不是跑完就结束。你需要知道每个 epoch 的 loss、验证集的准确率、当前的显存占用才能判断模型是否在正常学习。建议至少记录这些信息每个 epoch 的训练 loss 和验证 loss。学习率。显存占用情况。模型在验证集上的指标比如准确率。日志可以用简单的 print也可以用 TensorBoard 或 wandb。新手先不要引入太多工具print 足够你判断模型是否正常。模型保存有两种常见方式。一种是只保存参数torch.save(model.state_dict(), model_weights.pth)加载时model SimpleNet() model.load_state_dict(torch.load(model_weights.pth))另一种是保存整个模型torch.save(model, model_full.pth)推荐使用第一种因为只保存参数的兼容性更好模型结构可以通过代码重新构建不容易出现结构不一致的问题。保存模型时文件名里最好带上 epoch 和指标信息比如model_epoch5_acc0.92.pth。批量跑多个实验时这个习惯能省很多时间。5.3 低配置机器的参数调整如果你的电脑没有 NVIDIA 显卡或者显存只有几 GB并不代表不能学习。只要你做好参数调整。几个常见的调整思路减小 batch_size从默认的 32 降到 16 或 8。降低输入图片的分辨率。使用更小的模型比如减少卷积层通道数。使用 CPU 训练时把num_workers设为 0避免多进程加载导致的额外开销。训练时关闭不需要的计算比如验证阶段用torch.no_grad()。torch.no_grad()很实用with torch.no_grad(): for images, labels in val_loader: outputs model(images)它会让 PyTorch 不构建计算图省去梯度计算的内存和时间。验证和推理阶段都应该使用。6. 常见报错和排查链路6.1 启动阶段的报错启动阶段最常见的是 import 报错和 CUDA 相关的错误。ImportError: DLL load failed在 Windows 上经常出现。原因一般是 CUDA 和 PyTorch 版本不匹配或者缺少 Visual C Redistributable。建议先安装最新的 Visual C 运行库再重新安装 PyTorch。AssertionError: Torch not compiled with CUDA enabled表示你安装的是 CPU 版 PyTorch但代码里使用了 CUDA。解决办法是安装 GPU 版 PyTorch。ModuleNotFoundError: No module named torch这个最简单就是没安装或者没在当前 conda 环境里安装。检查一下当前环境的 Python 路径是不是你期望的那个。6.2 训练阶段的报错训练阶段最常见的报错是设备不一致RuntimeError: Expected all tensors to be on the same device模型参数在 GPU 上但输入数据在 CPU 上。解决办法是统一用to(device)。第二个常见问题是显存不足CUDA out of memory这时候不要急着加显存先按顺序排查batch_size 是不是太大输入分辨率是不是太高模型参数是否在逐步累积是否有变量被错误地保留而未释放。一个临时验证方法是把 batch_size 减半看是否能跑通。第三个常见问题是 loss 变成 NaN。一般的排查顺序是学习率过高、数据里包含 NaN 值、使用了不支持的精度格式。先检查数据和 label再降低学习率最后检查精度设置。6.3 排查顺序遇到问题我一般先按这个顺序排查看现象是报错、卡住、还是输出异常。看输入文件路径、文件格式、编码、数据里是否有空值。看环境Python 版本、PyTorch 版本、CUDA 版本、驱动版本。看参数batch_size、学习率、模型路径、输出目录。看工具本身这个功能是否在当前版本支持是否有已知限制。注意很多看似模型的问题其实是数据路径或权限问题。如果训练一开始就报错先检查文件是不是真的存在、路径里是不是包含中文或空格、所在目录能不能写文件。7. 一周实战计划每天做什么7.1 每天任务安排最后给一个可以直接参考的一周计划。第一天安装 Anaconda创建 Python 3.10 环境安装 CPU 版或 GPU 版 PyTorch。验证import torch和torch.cuda.is_available()。这一天的核心是环境跑通不要写任何模型代码。第二天学习张量基础操作把创建、索引、切片、reshape、设备移动都过一遍。再用 NumPy 和 Torch 做一次数据互转。第三天理解自动求导写一个简单函数的梯度计算再写一个线性回归模型。确认 loss 在下降。第四天学习nn.Module的写法实现一个含一层隐藏层的全连接网络在 MNIST 或模拟数据上跑通训练循环。第五天引入 Dataset 和 DataLoader把 MNIST 数据用标准方式加载加入验证集打印每个 epoch 的训练 loss 和验证准确率。第六天加入模型保存和加载实现从头训练的完整流程。用torch.no_grad()做推理验证。如果显存或内存不足尝试降低 batch_size 和分辨率。第七天了解 FP16、BF16、TF32 的差异在自己训练的模型上试一下混合精度训练。再把模型导出成 ONNX 或加载保存的权重做推理完成从训练到部署的最小闭环。7.2 学习中的误区第一个误区是花太多时间比较框架。PyTorch 和 TensorFlow 之争不是新手该操心的事先把一个框架用熟比反复切换框架有价值得多。第二个误区是跳过环境搭建直接复制代码。代码复制下来跑不动大概率是环境差异导致的。环境搭好比多写十行代码重要。第三个误区是盲目追求大模型、大 batch。新手一开始跑大模型只会卡在显存和内存上对理解原理没有帮助。先跑小模型把训练循环彻底理解再逐步扩大规模。第四个误区是只跑通一次就结束。跑通一次只能说明代码没有低级错误不代表你对训练流程有掌控力。试着改学习率、改 batch_size、加数据增强观察 loss 和指标的变化这才是真正的学习。PyTorch 的学习难度并不在 API 本身而在你是否能把环境、数据、模型、训练、保存、部署这几个环节串起来。按上面的顺序走一遍一周时间足够建立起一个完整的框架认知。之后再遇到更复杂的模型和任务你会发现核心思路是一样的准备好数据构建模型设计训练循环验证效果保存结果。剩下的只是在这个基础上做扩展。