2024年时候有个很有意思的现象身边原本写TensorFlow的老伙计一个个都悄悄把主力框架换成了PyTorch。不只是学术圈子工业界的招聘JD上也越来越频繁地出现“熟悉PyTorch优先”。我自己的学习路径是从TensorFlow 1.x时代摸过来的中间被静态图的调试折磨过不少次后来转PyTorch之后才真正体会到“动态图加自动求导”有多顺手。这篇笔记算是我把踩过的坑、看过的源码、实际项目的经验重新梳理了一遍写给想从零开始或者是被环境配置劝退过的朋友。文章覆盖范围比较广从环境搭建、张量和反向传播的原理到完整训练一个图像分类模型再到LSTM、注意力机制、模型转ONNX这些进阶话题按顺序读就行也可以直接跳到对应章节查答案。1. 劝退最多人的不是算法是环境搭建先说一个反直觉的事实很多人学PyTorch第一天就放弃了不是卡在张量乘法也不是卡在反向传播而是卡在了“装完环境之后根本 import 不进来”。PyTorch的环境配置之所以劝退是因为它牵扯到Python版本、CUDA驱动、显卡驱动、包管理器、操作系统好几条线任何一环不匹配就是各种奇奇怪怪的报错。1.1 Anaconda虚拟环境隔离是省心的前提我最开始学的时候直接在系统Python里pip install torch后来才发现这是给自己埋雷。不同项目的依赖是互相打架的你今天装了个需要Python 3.10的项目明天另一个项目又要求3.8系统环境很快就成了一锅粥。所以我建议的第一步永远是先装Anaconda或者Miniconda。conda create -n pytorch python3.10 conda activate pytorch创建虚拟环境之后再安装PyTorch就相对安全了。这里有个小建议conda镜像和pip镜像在国内环境下有时候会很慢建议给conda配好清华镜像源给pip配好清华PyPI镜像。具体配置方法网上随便一搜就有我在这里就不赘述了。1.2 GPU版本还是CPU版本先回答三个问题很多人上来就搜“pytorch安装教程gpu”但实际自己到底需不需要GPU其实没想清楚。我一般会问三个问题你的电脑有没有NVIDIA独立显卡没有的话直接装CPU版本别折腾。你的显卡显存多大跑个小模型、做做学习实验4GB以上就够入门了。你的CUDA驱动版本支持哪个CUDA toolkit这个用nvidia-smi命令能看到。如果显卡是AMD或者Intel的核显或者你是Mac用户那PyTorch也有对应的CPU版本甚至MPS加速支持。比如有个热词是“安装pytorch是不是必须装有GPU”这里可以明确回复不是。GPU只影响训练速度不影响你能不能学。用CPU版本跑MNIST、跑小型的全连接网络完全没问题只是跑大模型或者大batch的时候你会觉得时间走得特别慢而已。确定了这三个问题之后去PyTorch官网首页找下面这种命令就行# CUDA 12.4 版本 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 # CPU 版本 pip3 install torch torchvision torchaudio安装之前你最好确认一下python -c import torch; print(torch.__version__)能够正常输出再顺手print(torch.cuda.is_available())看看GPU是否可用。1.3 WSL2方案在Windows底下获得接近原生的体验热词里有“pytorch环境搭建wsl”和“7900xtx pytorch wsl”说明这块现在也有不少人关心。WSL2本身就是Windows Subsystem for LinuxWindows 10/11用户开启它之后可以原生跑一个Linux环境。这对PyTorch用户有几个好处很多开源项目只给了Linux的安装脚本你在WSL2里直接跑就行完全不用把代码搬到云服务器其次WSL2里Linux内核对于NVIDIA显卡的驱动支持现在很成熟torch.cuda.is_available()通常返回True。具体步骤大致是# 在管理员PowerShell里启用WSL2 wsl --install # 安装Ubuntu发行版 wsl --install -d Ubuntu-22.04进入Ubuntu终端之后先更新系统再装Miniconda接下来就完全照搬Linux环境下的安装流程了。需要注意的一点是WSL2里是没有/usr/local/cuda这种东西的它依赖Windows侧的NVIDIA驱动所以不要单独在WSL2里再装一套显卡驱动否则有时候会搞出兼容性问题。1.4 离线安装与特殊架构的兜底方案如果你所在的环境没法访问外网那就只能走离线安装路线。在能联网的机器上下载对应的wheel包然后拷贝过去用pip install xxx.whl离线安装。关键是选对版本比如torch-2.3.1cu121-cp310-cp310-linux_x86_64.whl这个文件名里面cp310表示Python 3.10cu121表示CUDA 12.1linux_x86_64表示Linux 64位任何一个不匹配都会安装失败。另一个特殊场景是国产化环境比如热词里提到的“麒麟系统v10 海光GPU安装pytorch”。海光GPU目前走的是ROCm生态PyTorch官方其实已经逐步支持ROCm了所以思路就是先确认系统里的ROCm版本然后到PyTorch官网找对应的ROCm wheel找不到的情况下再看厂商有没有提供适配包。总体原则就是提前确认三件事系统架构x86还是arm、Python版本、GPU的软件生态CUDA还是ROCm然后对着选包不要盲目复制网上的命令。2. 张量与自动求导先忘掉手写梯度那套东西环境搞定之后第一个要过的概念关就是张量。我见过很多初学者把PyTorch的Tensor当成了普通的数组来用其实它远远不止是个多维数组。2.1 张量带方向盘的多维数组你可以把Tensor理解成一辆装了方向盘和油门的车。普通数组你只能“看”里面的数据而Tensor这个东西不仅存了数据还记住了“我是怎么算出来的”。这句话是理解整个自动求导的关键。创建一个Tensor很简单import torch # 从列表创建 a torch.tensor([1, 2, 3.0]) # 全零 b torch.zeros(2, 3) # 随机初始化 c torch.randn(2, 3, requires_gradTrue) # 创建之后的数据类型和形状 print(a.dtype, a.shape) # torch.float32, torch.Size([3])这里有个很常见的坑torch.tensor([1, 2, 3])创建出来的是整数类型int64如果你后面想让它参与梯度计算可能会报错。所以我一般会写成torch.tensor([1.0, 2.0, 3.0])或者用tensor tensor.float()显式转换。再说一个偏远但实用的点热词里有个“pytorch返回实例的类对象名称”这个说的是在调试代码的时候你想知道一个模型或者层到底是什么类。常规做法是type(model)但更好的做法是model torch.nn.Linear(10, 5) print(model.__class__.__name__) # Linear print(type(model)) # class torch.nn.modules.linear.Linear在debug大型模型结构的时候用__class__.__name__拿到的纯字符串更干净打印结构图也方便这是我实际写项目时的一个小习惯。2.2 自动求导的底层逻辑计算图就是一个记账本PyTorch的自动求导机制官方叫Autograd底层是靠计算图实现的。我经常用记账本来类比你每次做一次张量运算PyTorch都会在背后记一笔账——“这个结果是由哪两个数怎么加出来的、乘出来的”。当你调用backward()的时候它就像把整本账从后往前翻一遍按照链式法则把每个变量的梯度算出来。x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # tensor([7.])上面这个例子里y x² 3x 1的导数是2x 3所以x2时梯度正好是7。backward()执行的就是反向传播算法它会自动沿着计算图回传把每个叶子节点的梯度算出来。如果你想让某个变量不参与梯度计算可以用requires_gradFalse或者with torch.no_grad():。推理阶段的循环里我强烈建议加上no_grad因为推理不需要梯度不关掉的话会白白浪费大量显存和算力。2.3 backward() 的真实行为和常见的坑backward()这个接口看着简单用起来还是有几个坑的。第一个坑是梯度累积。PyTorch默认每次backward()之后不会自动清零梯度而是把新算出来的梯度累加到上一次的数值上。这意味着如果你在训练循环里忘了optimizer.zero_grad()模型参数会以肉眼可见的速度失控损失函数直接变成NaN或者跳变。正确的训练循环模板一定是先清空梯度、再前向计算、再反向传播、最后更新参数。第二个坑是标量要求。backward()默认是要求目标是一个标量如果你对一个包含多个元素的张量直接调用backward()就会报错提示“grad can be implicitly created only for scalar outputs”。这种时候要么你先对张量求和变成一个标量要么传一个和形状相同的权重参数进去。第三个坑是数据类型和设备的匹配。CPU上的Tensor参与GPU上的Tensor运算一定会报错float32和float64混着算有时候也会出问题。这些坑单独看都很小但每一个都能让你debug一个晚上。经验就是先看报错信息的最后一行90%的autograd问题都在最后一行里写明白了。3. 搭建第一个训练任务用MNIST手写数字分类串起全流程学PyTorch最容易获得成就感的一个小项目就是用手写数字识别练手。模型简单数据集也不需要额外下载跑通整个流程之后你对“训练”这件事就有了整体感知。这里我选全连接神经网络而不是CNN因为初学者应该先把训练循环里的每个组件看明白而不是一开始就被卷积、池化这些层分散注意力。3.1 数据管道四件套Dataset、DataLoader、transform、normalizePyTorch里面数据这块由几个组件配合完成。torchvision.datasets.MNIST负责下载和管理数据文件transform负责对图片做预处理比如把PIL图像转成Tensor、做归一化DataLoader负责把数据打包成一个个batch并且支持多线程加载和乱序打乱。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( root./data, trainTrue, downloadTrue, transformtransform ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2)这里要详细说说归一化为什么重要。MNIST数据集的像素值范围是0到255如果不做归一化直接扔进模型那么第一层的输出会非常大激活函数很容易进入饱和区梯度要么消失要么爆炸。Normalize((0.1307,), (0.3081,))里面的两个数分别是MNIST数据集的均值和标准差它的作用是把数据分布拉到一个接近标准正态分布的状态。以后你处理自己的数据时也先花两分钟算一下均值和标准差再训练这是提升训练稳定性的一个关键细节。3.2 用nn.Module定义模型forward远比你想的重要nn.Module是PyTorch里所有模型和网络层的基类。定义一个模型本质上就是做两件事在__init__里面搭建你的层在forward里面规定数据怎么流过这些层。import torch.nn as nn class MLP(nn.Module): def __init__(self, input_dim784, hidden_dim128, num_classes10): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_dim, num_classes) def forward(self, x): # x: [batch_size, 1, 28, 28] x x.view(x.size(0), -1) # 展平成 [batch_size, 784] x self.fc1(x) x self.relu(x) x self.fc2(x) return x这里特别强调一下forward方法。nn.Module的__call__内部会先去执行一些底层的hook逻辑然后再调用你定义的forward所以你在外部直接调用model(x)而不是model.forward(x)。这两者在语义上有区别直接调forward会跳过某些模块行为比如Dropout在推理和训练模式下的切换逻辑。正确的做法永远是output model(x)。3.3 三个循环训练、验证、保存与加载数据管道有了模型定义了接下来就是最关键的三段式流程。我把训练循环的基本模板放在下面这个模板你在后续几乎所有PyTorch项目里都能见到建议直接背下来。import torch.optim as optim model MLP() optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() for epoch in range(5): model.train() for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in valid_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(fEpoch {epoch}, accuracy: {correct / total:.4f})训练循环里有两个初学者经常疑惑的点。第一个是model.train()和model.eval()。这两个方法并不是用来控制梯度计算的它们是切换模型中那些在训练和推理时行为不同的层比如Dropout和BatchNorm。你如果不切eval()推理阶段的预测结果会因为Dropout随机丢弃而波动明明同样的输入可能得到不同的输出。第二个是torch.max(outputs, 1)。outputs的维度是[batch_size, num_classes]torch.max的第二个参数表示沿着哪一维找最大值。找出来的predicted是每个样本预测的类别索引然后和真实标签比较就能算准确率了。模型的保存与加载是另一个高频操作。PyTorch有两种主流保存方式我建议你保存state_dict而不是整个模型# 保存 torch.save(model.state_dict(), mnist_mlp.pt) # 加载 model MLP() model.load_state_dict(torch.load(mnist_mlp.pt)) model.eval()保存state_dict的好处是跨设备兼容性好而且只存了权重参数文件体积小。如果你还想把优化器的状态、当前epoch、随机种子这些一起保存下来做断点续训可以把它们打包成字典一起存。3.4 训练中的高频问题排查手册跑通MNIST之后你会开始遇到各种问题这里我整理一份高频问题排查清单每一类都是我实际遇到过的。loss不降先看是不是学习率太大或者太小。学习率大容易震荡甚至发散学习率小则收敛极慢像是蜗牛爬。建议从1e-3起步损失不降就依次尝试1e-4和1e-2。然后确认数据归一化有没有做标签有没有对应正确。维度对不上这是最常见的报错类型。做法是在forward里手动print每一步的shape一行一行检查很快就能定位到是哪一层把形状整错了。还有个小技巧写代码时多写注释把每个张量的形状标在旁边例如# [B, 784]。显存溢出OOM先检查一下num_workers是不是开太多了然后看batch size是不是太大了。实在不行就把模型放到CPU上跑训练慢一点总比直接崩溃好。结果为NaN在损失函数前后打印一下数值看是不是输入包含了Inf或者NaN梯度爆炸也经常导致NaN可以试试加一个gradient clippingtorch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这个技巧在训练RNN、Transformer这类模型时特别常用我现在写训练循环的时候基本上都会顺手加上。4. 进阶实战从玩具模型到真实项目当你把MNIST跑通、把训练循环彻底理解之后PyTorch的大门才算真正打开。这一节我会挑几个进阶方向来聊这些方向分别对应热词里的LSTM源码、attention机制、模型转ONNX、以及强化学习的TD3代码。4.1 序列模型怎么学从读LSTM源码开始热词里有“pytorch lstm源码”说明很多人会卡在循环神经网络这块。LSTM对外接口看着复杂其实就是三个参数的事input_size是每个时间步输入特征的维度hidden_size是隐状态向量的维度num_layers是堆叠的层数。import torch.nn as nn lstm nn.LSTM(input_size50, hidden_size128, num_layers2, batch_firstTrue) # x: [batch, seq_len, input_size] x torch.randn(32, 20, 50) output, (h_n, c_n) lstm(x) print(output.shape) # [32, 20, 128]很多人不理解output和(h_n, c_n)之间的关系。output保存的是最后一个LSTM层在每个时间步的输出shape是[batch, seq_len, hidden_size]h_n是最后一层的最终隐藏状态shape是[num_layers, batch, hidden_size]c_n是细胞状态。如果你想取最后一个时间步的输出用于分类可以用output[:, -1, :]它和h_n[-1]是同一个东西这是一个很容易被绕晕的点。我实际用LSTM的体会是如果你要处理的是长序列LSTM的梯度仍然容易不稳定记得加上上一节提到的梯度裁剪。如果你要处理的是超长序列或者需要考虑并行计算那LSTM就不太合适了需要转向Transformer。4.2 注意力机制seq2seq decoder里最核心的一环热词里有“a generic attention module for a decoder in seq2seq pytorch”这是一个非常经典的注意力模块。简单讲seq2seq模型里编码器把输入序列压缩成一个语义向量但如果句子很长单个向量装不下全部信息。注意力机制的思路是解码器在生成当前位置输出的时候不要只依赖那个向量而是“回看”编码器每一步的隐状态给它们分配不同的权重重点关注与当前输出最相关的部分。import torch import torch.nn.functional as F class Attention(nn.Module): def __init__(self, hidden_size): super().__init__() self.W nn.Linear(hidden_size, hidden_size, biasFalse) def forward(self, decoder_hidden, encoder_outputs): # decoder_hidden: [batch, hidden] # encoder_outputs: [batch, seq_len, hidden] score torch.bmm(encoder_outputs, self.W(decoder_hidden).unsqueeze(2)).squeeze(2) attn_weight F.softmax(score, dim1) context torch.bmm(attn_weight.unsqueeze(1), encoder_outputs).squeeze(1) return context, attn_weight这里面bmm是批量矩阵乘法score计算每个编码器位置的得分经过softmax转成注意力权重再用权重把编码器输出加权求和得到上下文向量。这个模块虽然短但它是整个注意力机制的骨架。往细了说还有加性注意力、点积注意力、多头注意力这些变体但原理都是“求相似度、算权重、加权求和”把这个思路吃透再看Transformer源码就会轻松很多。4.3 PyTorch到ONNX模型部署的关键一步热词里“pytorch转onnx”出现的频率很高这确实是很多人在模型开发完后面对的第一个部署问题。ONNX全称是开放神经网络交换格式它相当于一个跨框架的通用格式。你把PyTorch模型转成ONNX之后可以对接ONNX Runtime、TensorRT甚至可以在网页端跑。转ONNX的操作非常简单model MLP() model.load_state_dict(torch.load(mnist_mlp.pt)) model.eval() dummy_input torch.randn(1, 1, 28, 28) torch.onnx.export( model, dummy_input, mnist_mlp.onnx, opset_version17, input_names[input], output_names[output] )导出时有几个关键点容易被忽略。第一个model.eval()必须调用否则BatchNorm或Dropout的行为会被固化成训练态。第二个dummy_input的shape必须和实际输入一致特别是batch维度ONNX导出之后计算图结构就固定了。第三个opset版本不是越高越好要看目标推理引擎支持到多少。导出之后怎么验证呢用ONNX Runtime跑一下比对输出和PyTorch原模型的输出是否一致。我之前遇到过动态shape的处理问题如果输入序列长度可变需要在export时设置dynamic_axes参数这样才能保证不同长度的输入都能被ONNX模型处理。4.4 从TD3代码谈强化学习代码的阅读方法热词里“td3代码pytorch”把强化学习也带进来了。TD3全称是Twin Delayed DDPG一个基于演员-评论家架构的强化学习算法。说实话第一次看TD3代码时我也挺晕的因为里面又包含Actor网络、Critic网络、目标网络、经验回放缓冲池好几套结构。我推荐的学习方法不是从头到尾一行行读而是先找到主训练循环通常在一个train_one_step函数里理解每一步在做什么选择一个动作、环境执行、存储经验、从经验池采样、更新Critic、定期更新Actor、软更新目标网络。把这条主线理清楚之后再回头一层层看各个网络的定义就会清晰得多。# TD3训练循环里的极简核心示意 def train_step(self): s, a, r, s_next, done self.replay_buffer.sample(256) # 用目标网络计算目标Q值 target_q self.target_critic(s_next, self.target_actor(s_next)) target_q r self.gamma * (1 - done) * target_q # 更新Critic loss F.mse_loss(self.critic(s, a), target_q.detach()) self.critic_optimizer.zero_grad() loss.backward() self.critic_optimizer.step()强化学习代码和深度学习代码最大的差异是它里面到处是detach()和“目标网络延迟更新”这种反直觉的设计不管在读代码还是调参时都容易迷失。但好消息是只要你前面的autograd和训练循环基础牢固至少能看明白每一步在干嘛。4.5 最后一个实用技巧善用模型结构打印最后分享一个我每次拿到新模型都会做的动作打印模型结构。PyTorch里直接print(model)就能看到每一层的名字和参数规模我用这个来核对模型是否搭对、参数量是否符合预期。热词里那个“pytorch返回实例的类对象名称”也与此相关当你写通用代码或者调试脚本时用model.__class__.__name__判断当前模块类型比比较字符串型的type输出更加稳定。我个人在实际操作中的体会是PyTorch的学习曲线其实可以拆成三段第一段是会跑通官方MNIST例程搞清楚训练循环里每个组件是干嘛的第二段是能够自己定义模型、改造数据管道、处理训练中出现的各种问题第三段就是看懂进阶代码比如LSTM源码、注意力模块、强化学习算法并具备自己写出来的能力。到了第三段PyTorch对你来说就不再是一个库而是一种思维方式了。希望这篇笔记能帮你少走几个我走过的弯路剩下的就靠你自己去跑代码、踩坑、再总结这个过程本身就是最好的学习方式。 PyTorch实战笔记：从环境搭建到模型部署
2024年时候有个很有意思的现象身边原本写TensorFlow的老伙计一个个都悄悄把主力框架换成了PyTorch。不只是学术圈子工业界的招聘JD上也越来越频繁地出现“熟悉PyTorch优先”。我自己的学习路径是从TensorFlow 1.x时代摸过来的中间被静态图的调试折磨过不少次后来转PyTorch之后才真正体会到“动态图加自动求导”有多顺手。这篇笔记算是我把踩过的坑、看过的源码、实际项目的经验重新梳理了一遍写给想从零开始或者是被环境配置劝退过的朋友。文章覆盖范围比较广从环境搭建、张量和反向传播的原理到完整训练一个图像分类模型再到LSTM、注意力机制、模型转ONNX这些进阶话题按顺序读就行也可以直接跳到对应章节查答案。1. 劝退最多人的不是算法是环境搭建先说一个反直觉的事实很多人学PyTorch第一天就放弃了不是卡在张量乘法也不是卡在反向传播而是卡在了“装完环境之后根本 import 不进来”。PyTorch的环境配置之所以劝退是因为它牵扯到Python版本、CUDA驱动、显卡驱动、包管理器、操作系统好几条线任何一环不匹配就是各种奇奇怪怪的报错。1.1 Anaconda虚拟环境隔离是省心的前提我最开始学的时候直接在系统Python里pip install torch后来才发现这是给自己埋雷。不同项目的依赖是互相打架的你今天装了个需要Python 3.10的项目明天另一个项目又要求3.8系统环境很快就成了一锅粥。所以我建议的第一步永远是先装Anaconda或者Miniconda。conda create -n pytorch python3.10 conda activate pytorch创建虚拟环境之后再安装PyTorch就相对安全了。这里有个小建议conda镜像和pip镜像在国内环境下有时候会很慢建议给conda配好清华镜像源给pip配好清华PyPI镜像。具体配置方法网上随便一搜就有我在这里就不赘述了。1.2 GPU版本还是CPU版本先回答三个问题很多人上来就搜“pytorch安装教程gpu”但实际自己到底需不需要GPU其实没想清楚。我一般会问三个问题你的电脑有没有NVIDIA独立显卡没有的话直接装CPU版本别折腾。你的显卡显存多大跑个小模型、做做学习实验4GB以上就够入门了。你的CUDA驱动版本支持哪个CUDA toolkit这个用nvidia-smi命令能看到。如果显卡是AMD或者Intel的核显或者你是Mac用户那PyTorch也有对应的CPU版本甚至MPS加速支持。比如有个热词是“安装pytorch是不是必须装有GPU”这里可以明确回复不是。GPU只影响训练速度不影响你能不能学。用CPU版本跑MNIST、跑小型的全连接网络完全没问题只是跑大模型或者大batch的时候你会觉得时间走得特别慢而已。确定了这三个问题之后去PyTorch官网首页找下面这种命令就行# CUDA 12.4 版本 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 # CPU 版本 pip3 install torch torchvision torchaudio安装之前你最好确认一下python -c import torch; print(torch.__version__)能够正常输出再顺手print(torch.cuda.is_available())看看GPU是否可用。1.3 WSL2方案在Windows底下获得接近原生的体验热词里有“pytorch环境搭建wsl”和“7900xtx pytorch wsl”说明这块现在也有不少人关心。WSL2本身就是Windows Subsystem for LinuxWindows 10/11用户开启它之后可以原生跑一个Linux环境。这对PyTorch用户有几个好处很多开源项目只给了Linux的安装脚本你在WSL2里直接跑就行完全不用把代码搬到云服务器其次WSL2里Linux内核对于NVIDIA显卡的驱动支持现在很成熟torch.cuda.is_available()通常返回True。具体步骤大致是# 在管理员PowerShell里启用WSL2 wsl --install # 安装Ubuntu发行版 wsl --install -d Ubuntu-22.04进入Ubuntu终端之后先更新系统再装Miniconda接下来就完全照搬Linux环境下的安装流程了。需要注意的一点是WSL2里是没有/usr/local/cuda这种东西的它依赖Windows侧的NVIDIA驱动所以不要单独在WSL2里再装一套显卡驱动否则有时候会搞出兼容性问题。1.4 离线安装与特殊架构的兜底方案如果你所在的环境没法访问外网那就只能走离线安装路线。在能联网的机器上下载对应的wheel包然后拷贝过去用pip install xxx.whl离线安装。关键是选对版本比如torch-2.3.1cu121-cp310-cp310-linux_x86_64.whl这个文件名里面cp310表示Python 3.10cu121表示CUDA 12.1linux_x86_64表示Linux 64位任何一个不匹配都会安装失败。另一个特殊场景是国产化环境比如热词里提到的“麒麟系统v10 海光GPU安装pytorch”。海光GPU目前走的是ROCm生态PyTorch官方其实已经逐步支持ROCm了所以思路就是先确认系统里的ROCm版本然后到PyTorch官网找对应的ROCm wheel找不到的情况下再看厂商有没有提供适配包。总体原则就是提前确认三件事系统架构x86还是arm、Python版本、GPU的软件生态CUDA还是ROCm然后对着选包不要盲目复制网上的命令。2. 张量与自动求导先忘掉手写梯度那套东西环境搞定之后第一个要过的概念关就是张量。我见过很多初学者把PyTorch的Tensor当成了普通的数组来用其实它远远不止是个多维数组。2.1 张量带方向盘的多维数组你可以把Tensor理解成一辆装了方向盘和油门的车。普通数组你只能“看”里面的数据而Tensor这个东西不仅存了数据还记住了“我是怎么算出来的”。这句话是理解整个自动求导的关键。创建一个Tensor很简单import torch # 从列表创建 a torch.tensor([1, 2, 3.0]) # 全零 b torch.zeros(2, 3) # 随机初始化 c torch.randn(2, 3, requires_gradTrue) # 创建之后的数据类型和形状 print(a.dtype, a.shape) # torch.float32, torch.Size([3])这里有个很常见的坑torch.tensor([1, 2, 3])创建出来的是整数类型int64如果你后面想让它参与梯度计算可能会报错。所以我一般会写成torch.tensor([1.0, 2.0, 3.0])或者用tensor tensor.float()显式转换。再说一个偏远但实用的点热词里有个“pytorch返回实例的类对象名称”这个说的是在调试代码的时候你想知道一个模型或者层到底是什么类。常规做法是type(model)但更好的做法是model torch.nn.Linear(10, 5) print(model.__class__.__name__) # Linear print(type(model)) # class torch.nn.modules.linear.Linear在debug大型模型结构的时候用__class__.__name__拿到的纯字符串更干净打印结构图也方便这是我实际写项目时的一个小习惯。2.2 自动求导的底层逻辑计算图就是一个记账本PyTorch的自动求导机制官方叫Autograd底层是靠计算图实现的。我经常用记账本来类比你每次做一次张量运算PyTorch都会在背后记一笔账——“这个结果是由哪两个数怎么加出来的、乘出来的”。当你调用backward()的时候它就像把整本账从后往前翻一遍按照链式法则把每个变量的梯度算出来。x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # tensor([7.])上面这个例子里y x² 3x 1的导数是2x 3所以x2时梯度正好是7。backward()执行的就是反向传播算法它会自动沿着计算图回传把每个叶子节点的梯度算出来。如果你想让某个变量不参与梯度计算可以用requires_gradFalse或者with torch.no_grad():。推理阶段的循环里我强烈建议加上no_grad因为推理不需要梯度不关掉的话会白白浪费大量显存和算力。2.3 backward() 的真实行为和常见的坑backward()这个接口看着简单用起来还是有几个坑的。第一个坑是梯度累积。PyTorch默认每次backward()之后不会自动清零梯度而是把新算出来的梯度累加到上一次的数值上。这意味着如果你在训练循环里忘了optimizer.zero_grad()模型参数会以肉眼可见的速度失控损失函数直接变成NaN或者跳变。正确的训练循环模板一定是先清空梯度、再前向计算、再反向传播、最后更新参数。第二个坑是标量要求。backward()默认是要求目标是一个标量如果你对一个包含多个元素的张量直接调用backward()就会报错提示“grad can be implicitly created only for scalar outputs”。这种时候要么你先对张量求和变成一个标量要么传一个和形状相同的权重参数进去。第三个坑是数据类型和设备的匹配。CPU上的Tensor参与GPU上的Tensor运算一定会报错float32和float64混着算有时候也会出问题。这些坑单独看都很小但每一个都能让你debug一个晚上。经验就是先看报错信息的最后一行90%的autograd问题都在最后一行里写明白了。3. 搭建第一个训练任务用MNIST手写数字分类串起全流程学PyTorch最容易获得成就感的一个小项目就是用手写数字识别练手。模型简单数据集也不需要额外下载跑通整个流程之后你对“训练”这件事就有了整体感知。这里我选全连接神经网络而不是CNN因为初学者应该先把训练循环里的每个组件看明白而不是一开始就被卷积、池化这些层分散注意力。3.1 数据管道四件套Dataset、DataLoader、transform、normalizePyTorch里面数据这块由几个组件配合完成。torchvision.datasets.MNIST负责下载和管理数据文件transform负责对图片做预处理比如把PIL图像转成Tensor、做归一化DataLoader负责把数据打包成一个个batch并且支持多线程加载和乱序打乱。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( root./data, trainTrue, downloadTrue, transformtransform ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2)这里要详细说说归一化为什么重要。MNIST数据集的像素值范围是0到255如果不做归一化直接扔进模型那么第一层的输出会非常大激活函数很容易进入饱和区梯度要么消失要么爆炸。Normalize((0.1307,), (0.3081,))里面的两个数分别是MNIST数据集的均值和标准差它的作用是把数据分布拉到一个接近标准正态分布的状态。以后你处理自己的数据时也先花两分钟算一下均值和标准差再训练这是提升训练稳定性的一个关键细节。3.2 用nn.Module定义模型forward远比你想的重要nn.Module是PyTorch里所有模型和网络层的基类。定义一个模型本质上就是做两件事在__init__里面搭建你的层在forward里面规定数据怎么流过这些层。import torch.nn as nn class MLP(nn.Module): def __init__(self, input_dim784, hidden_dim128, num_classes10): super().__init__() self.fc1 nn.Linear(input_dim, hidden_dim) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_dim, num_classes) def forward(self, x): # x: [batch_size, 1, 28, 28] x x.view(x.size(0), -1) # 展平成 [batch_size, 784] x self.fc1(x) x self.relu(x) x self.fc2(x) return x这里特别强调一下forward方法。nn.Module的__call__内部会先去执行一些底层的hook逻辑然后再调用你定义的forward所以你在外部直接调用model(x)而不是model.forward(x)。这两者在语义上有区别直接调forward会跳过某些模块行为比如Dropout在推理和训练模式下的切换逻辑。正确的做法永远是output model(x)。3.3 三个循环训练、验证、保存与加载数据管道有了模型定义了接下来就是最关键的三段式流程。我把训练循环的基本模板放在下面这个模板你在后续几乎所有PyTorch项目里都能见到建议直接背下来。import torch.optim as optim model MLP() optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() for epoch in range(5): model.train() for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in valid_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(fEpoch {epoch}, accuracy: {correct / total:.4f})训练循环里有两个初学者经常疑惑的点。第一个是model.train()和model.eval()。这两个方法并不是用来控制梯度计算的它们是切换模型中那些在训练和推理时行为不同的层比如Dropout和BatchNorm。你如果不切eval()推理阶段的预测结果会因为Dropout随机丢弃而波动明明同样的输入可能得到不同的输出。第二个是torch.max(outputs, 1)。outputs的维度是[batch_size, num_classes]torch.max的第二个参数表示沿着哪一维找最大值。找出来的predicted是每个样本预测的类别索引然后和真实标签比较就能算准确率了。模型的保存与加载是另一个高频操作。PyTorch有两种主流保存方式我建议你保存state_dict而不是整个模型# 保存 torch.save(model.state_dict(), mnist_mlp.pt) # 加载 model MLP() model.load_state_dict(torch.load(mnist_mlp.pt)) model.eval()保存state_dict的好处是跨设备兼容性好而且只存了权重参数文件体积小。如果你还想把优化器的状态、当前epoch、随机种子这些一起保存下来做断点续训可以把它们打包成字典一起存。3.4 训练中的高频问题排查手册跑通MNIST之后你会开始遇到各种问题这里我整理一份高频问题排查清单每一类都是我实际遇到过的。loss不降先看是不是学习率太大或者太小。学习率大容易震荡甚至发散学习率小则收敛极慢像是蜗牛爬。建议从1e-3起步损失不降就依次尝试1e-4和1e-2。然后确认数据归一化有没有做标签有没有对应正确。维度对不上这是最常见的报错类型。做法是在forward里手动print每一步的shape一行一行检查很快就能定位到是哪一层把形状整错了。还有个小技巧写代码时多写注释把每个张量的形状标在旁边例如# [B, 784]。显存溢出OOM先检查一下num_workers是不是开太多了然后看batch size是不是太大了。实在不行就把模型放到CPU上跑训练慢一点总比直接崩溃好。结果为NaN在损失函数前后打印一下数值看是不是输入包含了Inf或者NaN梯度爆炸也经常导致NaN可以试试加一个gradient clippingtorch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这个技巧在训练RNN、Transformer这类模型时特别常用我现在写训练循环的时候基本上都会顺手加上。4. 进阶实战从玩具模型到真实项目当你把MNIST跑通、把训练循环彻底理解之后PyTorch的大门才算真正打开。这一节我会挑几个进阶方向来聊这些方向分别对应热词里的LSTM源码、attention机制、模型转ONNX、以及强化学习的TD3代码。4.1 序列模型怎么学从读LSTM源码开始热词里有“pytorch lstm源码”说明很多人会卡在循环神经网络这块。LSTM对外接口看着复杂其实就是三个参数的事input_size是每个时间步输入特征的维度hidden_size是隐状态向量的维度num_layers是堆叠的层数。import torch.nn as nn lstm nn.LSTM(input_size50, hidden_size128, num_layers2, batch_firstTrue) # x: [batch, seq_len, input_size] x torch.randn(32, 20, 50) output, (h_n, c_n) lstm(x) print(output.shape) # [32, 20, 128]很多人不理解output和(h_n, c_n)之间的关系。output保存的是最后一个LSTM层在每个时间步的输出shape是[batch, seq_len, hidden_size]h_n是最后一层的最终隐藏状态shape是[num_layers, batch, hidden_size]c_n是细胞状态。如果你想取最后一个时间步的输出用于分类可以用output[:, -1, :]它和h_n[-1]是同一个东西这是一个很容易被绕晕的点。我实际用LSTM的体会是如果你要处理的是长序列LSTM的梯度仍然容易不稳定记得加上上一节提到的梯度裁剪。如果你要处理的是超长序列或者需要考虑并行计算那LSTM就不太合适了需要转向Transformer。4.2 注意力机制seq2seq decoder里最核心的一环热词里有“a generic attention module for a decoder in seq2seq pytorch”这是一个非常经典的注意力模块。简单讲seq2seq模型里编码器把输入序列压缩成一个语义向量但如果句子很长单个向量装不下全部信息。注意力机制的思路是解码器在生成当前位置输出的时候不要只依赖那个向量而是“回看”编码器每一步的隐状态给它们分配不同的权重重点关注与当前输出最相关的部分。import torch import torch.nn.functional as F class Attention(nn.Module): def __init__(self, hidden_size): super().__init__() self.W nn.Linear(hidden_size, hidden_size, biasFalse) def forward(self, decoder_hidden, encoder_outputs): # decoder_hidden: [batch, hidden] # encoder_outputs: [batch, seq_len, hidden] score torch.bmm(encoder_outputs, self.W(decoder_hidden).unsqueeze(2)).squeeze(2) attn_weight F.softmax(score, dim1) context torch.bmm(attn_weight.unsqueeze(1), encoder_outputs).squeeze(1) return context, attn_weight这里面bmm是批量矩阵乘法score计算每个编码器位置的得分经过softmax转成注意力权重再用权重把编码器输出加权求和得到上下文向量。这个模块虽然短但它是整个注意力机制的骨架。往细了说还有加性注意力、点积注意力、多头注意力这些变体但原理都是“求相似度、算权重、加权求和”把这个思路吃透再看Transformer源码就会轻松很多。4.3 PyTorch到ONNX模型部署的关键一步热词里“pytorch转onnx”出现的频率很高这确实是很多人在模型开发完后面对的第一个部署问题。ONNX全称是开放神经网络交换格式它相当于一个跨框架的通用格式。你把PyTorch模型转成ONNX之后可以对接ONNX Runtime、TensorRT甚至可以在网页端跑。转ONNX的操作非常简单model MLP() model.load_state_dict(torch.load(mnist_mlp.pt)) model.eval() dummy_input torch.randn(1, 1, 28, 28) torch.onnx.export( model, dummy_input, mnist_mlp.onnx, opset_version17, input_names[input], output_names[output] )导出时有几个关键点容易被忽略。第一个model.eval()必须调用否则BatchNorm或Dropout的行为会被固化成训练态。第二个dummy_input的shape必须和实际输入一致特别是batch维度ONNX导出之后计算图结构就固定了。第三个opset版本不是越高越好要看目标推理引擎支持到多少。导出之后怎么验证呢用ONNX Runtime跑一下比对输出和PyTorch原模型的输出是否一致。我之前遇到过动态shape的处理问题如果输入序列长度可变需要在export时设置dynamic_axes参数这样才能保证不同长度的输入都能被ONNX模型处理。4.4 从TD3代码谈强化学习代码的阅读方法热词里“td3代码pytorch”把强化学习也带进来了。TD3全称是Twin Delayed DDPG一个基于演员-评论家架构的强化学习算法。说实话第一次看TD3代码时我也挺晕的因为里面又包含Actor网络、Critic网络、目标网络、经验回放缓冲池好几套结构。我推荐的学习方法不是从头到尾一行行读而是先找到主训练循环通常在一个train_one_step函数里理解每一步在做什么选择一个动作、环境执行、存储经验、从经验池采样、更新Critic、定期更新Actor、软更新目标网络。把这条主线理清楚之后再回头一层层看各个网络的定义就会清晰得多。# TD3训练循环里的极简核心示意 def train_step(self): s, a, r, s_next, done self.replay_buffer.sample(256) # 用目标网络计算目标Q值 target_q self.target_critic(s_next, self.target_actor(s_next)) target_q r self.gamma * (1 - done) * target_q # 更新Critic loss F.mse_loss(self.critic(s, a), target_q.detach()) self.critic_optimizer.zero_grad() loss.backward() self.critic_optimizer.step()强化学习代码和深度学习代码最大的差异是它里面到处是detach()和“目标网络延迟更新”这种反直觉的设计不管在读代码还是调参时都容易迷失。但好消息是只要你前面的autograd和训练循环基础牢固至少能看明白每一步在干嘛。4.5 最后一个实用技巧善用模型结构打印最后分享一个我每次拿到新模型都会做的动作打印模型结构。PyTorch里直接print(model)就能看到每一层的名字和参数规模我用这个来核对模型是否搭对、参数量是否符合预期。热词里那个“pytorch返回实例的类对象名称”也与此相关当你写通用代码或者调试脚本时用model.__class__.__name__判断当前模块类型比比较字符串型的type输出更加稳定。我个人在实际操作中的体会是PyTorch的学习曲线其实可以拆成三段第一段是会跑通官方MNIST例程搞清楚训练循环里每个组件是干嘛的第二段是能够自己定义模型、改造数据管道、处理训练中出现的各种问题第三段就是看懂进阶代码比如LSTM源码、注意力模块、强化学习算法并具备自己写出来的能力。到了第三段PyTorch对你来说就不再是一个库而是一种思维方式了。希望这篇笔记能帮你少走几个我走过的弯路剩下的就靠你自己去跑代码、踩坑、再总结这个过程本身就是最好的学习方式。