深度学习中,PyTorch这个名字你大概率绕不开。不管你是准备入行的新手,还是做过几个月项目的老手,只要打开任何一个开源项目、竞赛代码或者论文附带实现,几乎都是用PyTorch写的。它可以用来做图像识别、目标检测、自然语言处理、语音识别,甚至强化学习,基本上你听过的主流AI应用,都能用这套框架跑起来。这篇文章不打算写成官方文档式的教程,就按我自己的上手经历,把PyTorch从环境搭建、核心概念到一个能跑的图像分类项目,再到训练过程中那些让人头疼的报错,完整地捋一遍。不管你完全没接触过,还是已经踩过几个坑,应该都能找到有用的东西。
1. 先搞懂PyTorch到底是什么,以及它的3个底层逻辑
很多初学者一上来就装环境、跑代码,结果代码是跑通了,但换个任务就完全懵了。我建议你先花十几分钟把PyTorch最核心的3个底层逻辑搞清楚,后面会省非常多的事。
1.1 张量(Tensor)不是普通的数组
一句话来说,张量就是带了两样特殊能力的多维数组。它和Python里的列表、NumPy里的数组长得像,但它多出来两个杀手锏:
第一个是GPU加速。普通的数组只能放在CPU上计算,而张量可以通过.to("cuda")直接搬到显卡上,并行计算的速度能差几十倍。你后面训练神经网络,几乎所有的时间都在做矩阵运算,没有GPU加速,一个稍微像样点的模型可能要跑几天几夜。
第二个是自动求导。这是神经网络训练的地基。一个模型要学习,就需要算出损失函数对每个参数的梯度,然后根据梯度去更新参数。在没有自动求导之前,你得自己用数学公式手推梯度,复杂网络根本推不动。PyTorch张量只要设了requires_grad=True,它在参与运算后,会自动把计算过程记录下来,你调用.backward(),梯度就自动算好了。
看个最朴素的例子:
import torch x = torch.tensor(2.0, requires_grad=True) y = x ** 2 y.backward() print(x.grad) # 输出 tensor(4.)这背后其实就是微积分里的链式法则:y对x求导等于2x,x等于2,所以结果是4。PyTorch帮你把这件事自动化了,而神经网络里那些动辄几百万个参数的求导,本质上也就是这样一层一层链式算下来的。
1.2 动态计算图是PyTorch最让人舒服的设计
计算图在深度学习里就是记录数据怎么经过各种运算最后得到结果的路线图。PyTorch用的是动态计算图,意思是图在你每执行一行代码时即时构建、即时运行。
这个设计带来的最大好处是调试极其直观。你可以随时打印中间变量的值,随时打断,甚至可以像写普通Python程序一样用if、for循环。很多其他框架用的是静态计算图,得先把整张图定义好,再喂数据进去运行,中间哪个步骤出错了,排查起来非常绕,就像你先写完整本菜谱再进厨房实际操作,写完才发现在第二页就有个菜谱步骤写错了,血压直接拉满。
PyTorch这种“运行的时候边算边记录”的方式,对于研究、实验、快速迭代来说简直是降维打击。你在开发时写起来更像是在写数学推导过程,而不是在伺候一个框架的语法。
1.3 你只需要这4个基础,就能开始学
不少人在学PyTorch前总觉得自己“基础不够”。根据我带过几个朋友上手的经验,你只需要4样东西:
- Python基础:列表、字典、函数、类、
for循环,够用就行 - NumPy一点的观念:知道数组有shape(维度形状)、有索引,这个概念对理解张量非常有帮助
- 机器学习最基本的感觉:知道什么是训练集、测试集,知道模型是拿数据学出来的,而不是提前写死规则的
- 一点线性代数直觉:知道矩阵乘法、维度匹配就完全够了
没有的话也不用怕,只要你在跑第一个项目时,愿意把每行代码执行后张量的形状打印出来看一眼,基础会在实战中不知不觉补齐。
2. 手把手把环境装好:显卡、CUDA和PyTorch版本怎么匹配
环境配不好,后面跑什么都是报错。但这一关真的不难,关键是把匹配逻辑搞清楚。
2.1 先搞明白CUDA版本这件事
很多人在装PyTorch时最纠结的就是CUDA版本。我帮你把逻辑理清楚,一共就三层:
- 你的显卡决定能不能用GPU加速
- 你的显卡驱动决定支持的最高CUDA版本
- 你装的PyTorch里捆绑了它需要的CUDA运行库
这里有个很关键的认知:PyTorch要求的CUDA版本不用和显卡驱动支持的CUDA版本完全相等,只要PyTorch要求的版本不超过驱动支持的版本就行。你可以把驱动支持的CUDA看成是容器能装多满的水,而PyTorch需要的只是其中一瓢。
在命令行输入下面这行,就能看到你的显卡驱动支持的最高CUDA版本:
nvidia-smi输出信息里右上角的CUDA Version,就是驱动支持的最高版本。如果这里显示的是12.x,那你可以放心装PyTorch官方提供的CUDA 11.8或12.1版本的包,都没问题。
2.2 用conda建独立环境是最省心的做法
我见过太多人在系统自带的Python里直接干,结果装包时要么提示没有权限,要么把项目A的依赖升级了导致项目B跑不起来。推荐的做法是给每个项目建一个独立的Python环境,隔离依赖,互不污染。
如果你装了Anaconda或者Miniconda,打开命令行执行:
conda create -n pytorch python=3.9 -y conda activate pytorch如果你不想装conda,用Python自带的venv也一样,原理相通。我个人的建议是新手用conda,因为后面装PyTorch、装CUDA相关组件时,conda帮你处理依赖问题会省很多事,而且你在Windows下遇到的环境问题往往更少。
2.3 安装命令怎么选
最推荐的方式是打开PyTorch官网,选择你的操作系统、包管理器、CUDA版本,它会自动生成对应的命令。不过我也总结了一张速查表,覆盖大部分常见情况:
| 你的场景 | 安装命令 |
|---|---|
| 有NVIDIA显卡,驱动支持CUDA 12.x | pip install torch torchvision torchaudio |
| 有NVIDIA显卡,驱动只支持到CUDA 11.8 | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 |
| 没有NVIDIA显卡,只能用CPU | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu |
| 安装anaconda后想用conda装GPU版 | conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia |
装完之后,用一个简短的脚本验证:
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,九成是装的CPU版,或者显卡驱动版本过低,按照上面的表格换命令重装一次就好。
注意:装完PyTorch后,不要随手
pip install torch再装一遍,那样很可能会把GPU版覆盖成CPU版,或者把版本搞乱。检查环境就认准pip list | grep torch。
3. 从零实现一个手写数字识别:完整项目实战
环境装好了,接下来就是动手跑通一个项目。我强烈推荐从手写数字识别开始,因为这可能是深度学习里最经典、最简单、数据量又足够小的入门任务。你的目标很简单:给一张28x28像素的手写数字灰度图,模型要判断它是0到9中的哪一个。
3.1 数据准备:Dataset和DataLoader是怎么回事
PyTorch里准备数据基本就靠两个类:Dataset负责“读取每一条数据”,DataLoader负责“批量、打乱、多线程地把数据喂给模型”。
对于手写数字这个任务,直接用现成的数据集就行:
from torchvision import datasets, transforms transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root='./data', train=False, download=True, transform=transform ) from torch.utils.data import DataLoader train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)有几个点我来解释一下为什么这么写:
ToTensor()把PIL图像转成张量,并且把像素值从0到255缩放到0到1之间,这是为了让神经网络的计算更稳定。Normalize((0.1307,), (0.3081,))是MNIST数据集的均值和标准差。为什么要做这步?因为很多算法在输入数据均值接近0、波动范围适中的情况下收敛更快。相当于先把数据“居中”再“归一”一遍。batch_size=64表示每次给模型64张图一起算。为什么不一次喂全部数据?因为一次全喂显存会爆,而且梯度计算时用一小批数据的平均梯度去更新参数,已经有很不错的稳定性和计算效率。shuffle=True只在训练集上开,目的是让每个batch里数据的顺序随机化,避免模型学到数据排列里的规律。测试集不需要。num_workers=2表示用2个子进程提前加载数据,减少训练时等数据的空闲时间。
3.2 定义网络结构:nn.Module的使用套路
PyTorch里定义神经网络,统一套路就是继承nn.Module,在__init__里定义网络层,在forward里定义数据怎么从输入走到输出。
import torch.nn as nn import torch.nn.functional as F class Net(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1) self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) return F.log_softmax(x, dim=1)这个网络结构很经典:两个卷积层提取图像特征,最大池化压缩特征尺寸,全连接层做分类,中间用Dropout防止过拟合。
有几个新手容易忽视的细节:
- 卷积核是3x3、步长是1、没有padding,所以
28x28的输入经过一次卷积变成26x26,再经过一次变成24x24,池化后变成12x12。第二个卷积层输出64个通道,所以进入全连接层之前,展平后的维度是64*12*12=9216。这就是fc1 = nn.Linear(9216, 128)里的9216是怎么来的。 forward函数定义了前向传播过程,而反向传播不需要你管,autograd会自动根据你写的前向过程算好梯度。这就是PyTorch最舒服的地方。
3.3 训练循环:最核心的4行代码
训练过程本质上是一个不断重复的循环:拿数据、算预测、算损失、算梯度、更新参数。用代码写出来就是:
import torch.optim as optim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Net().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) # 这4行是训练的核心,顺序一个都不能乱 optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print(f"第{epoch}轮, 第{batch_idx}批, 损失: {loss.item():.6f}")我把这4行掰开揉碎讲一下:
optimizer.zero_grad():把上一次迭代留下的梯度清零。因为PyTorch的backward()会把梯度累加到参数的.grad属性上,不清零的话,上一轮的梯度会和这一轮的混在一起,参数更新方向就错了。output = model(data):前向传播,也就是让数据通过网络算出一个预测结果。loss.backward():反向传播,自动算出每个参数相对于损失的梯度。optimizer.step():拿着算好的梯度,按照优化器的规则去更新参数。比如Adam就按它自己的一套自适应学习率逻辑来更新。
顺序能不能乱?不能。你如果没有先zero_grad就去backward,梯度会叠加;如果你step了才backward,那相当于用的上上次的梯度更新参数。这几个动作环环相扣,是深度学习中最重要的“肌肉记忆”。
3.4 模型评估与保存:别等训完了才想起来写
训练过程里,每训练几轮就应该在验证集上评估一次,看看模型是变好了还是在过拟合。评估代码和训练代码有区别,关键点在于不需要计算梯度:
def test(model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) accuracy = correct / len(test_loader.dataset) print(f"测试集平均损失: {test_loss:.4f}, 准确率: {accuracy:.4f}")这里有几个很容易踩的坑:
model.eval()是必须的。它会告诉模型现在是评估阶段,像Dropout和BatchNorm这类在训练和推理时行为不一样的层要切换状态。如果你忘了调用,用训练模式跑测试,Dropout还在随机丢神经元,结果每次预测都不稳定,准确率会忽高忽低。with torch.no_grad():是评估的标配。它告诉PyTorch这整段代码不需要追踪梯度,能省下大量显存和计算时间。推理阶段根本不需要梯度,开这个是一种好习惯。model.eval()只是切换模式,不代表参数不会再变。如果你想继续训练,下一轮记得调回model.train()。
模型训练完,保存的时候我推荐保存state_dict,也就是模型的参数字典,而不是整个模型对象:
torch.save(model.state_dict(), "mnist_cnn.pt")加载的时候也要先重建一个相同结构的Net(),再把参数填进去:
model2 = Net().to(device) model2.load_state_dict(torch.load("mnist_cnn.pt")) model2.eval()保存state_dict的好处是文件更小、更安全,而且是业界最通用的做法。加载后记得调用.eval(),不然新加载的模型还停在训练状态。
4. 训练神经网络最容易踩的5类坑:排查实录与避坑技巧
这部分就是干货中的干货了。我在自己写代码、帮别人调代码的过程中,归纳出下面这几类出现频率最高的报错和问题,你大概率也会遇到。
4.1 张量维度对不上:最高频的报错
这类报错最常见的长这样:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x128 and 9216x128)原因很简单:你在某一层把数据的维度搞错了。尤其容易发生在从卷积层跳到全连接层的那个地方,忘了flatten,或者在forward里修改了网络结构、忘记了同步修改fc1的输入维度。
我排查维度问题最快的方法是:在forward里多打印几次形状。
def forward(self, x): print("输入:", x.shape) x = self.conv1(x) print("conv1后:", x.shape) x = F.max_pool2d(x, 2) print("池化后:", x.shape) x = torch.flatten(x, 1) print("展平后:", x.shape) ...跑一次前向传播,立刻就能定位形状是在哪里对不上的。找到之后,把打印删掉,修正匹配的维度就行。
4.2 GPU显存不足(OOM)与CUDA不可用
“显存不足”是训练中另一个高频问题,报错通常长这样:
CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 6.00 GiB total capacity)注意一个容易迷惑的细节:它说Tried to allocate 20.00 MiB,并不是说你的模型只缺这20MB,而是说显卡到这一步确实剩不下更多空间了。排查思路按优先级来:
- 减小
batch_size:从64降到32、16,一般都能立刻解决问题 - 检查是不是有大量的张量被
to(device)搬到了GPU上,用完没释放 - 看看训练代码里的
with torch.no_grad()有没有忘加,不需要梯度的中间结果也会占显存 - 调用
torch.cuda.empty_cache()释放PyTorch缓存的显存碎片
如果torch.cuda.is_available()返回False,则先检查nvidia-smi的CUDA版本能不能支撑你装的PyTorch;再检查你是不是装在CPU版上了;最后看看驱动要不要更新。
4.3 数据加载卡顿与过拟合处理
数据加载慢的表现是:训练时GPU使用率忽高忽低,nvidia-smi里GPU利用率只有百分之几,CPU跑满。这通常意味着数据加载成了瓶颈。解决办法是把num_workers调大,比如设成4或8,让更多进程来预处理数据。
不过这也不是越大越好,num_workers超过CPU核心数时反而会增加进程切换开销,在Windows上num_workers设置过高还可能直接报异常报错,如果遇到莫名崩溃,把它设成0试试,0表示使用主进程加载数据,最稳但慢。
过拟合的表现是:训练集准确率很高,验证集准确率一直上不去。处理手段无非这几板斧:加Dropout、加数据增强(比如旋转、平移、加噪声)、降低模型复杂度、加早停(验证集准确率连续几轮不涨就停止训练)。新手最容易忽略的是在训练集上也能看出来过拟合的苗头:如果训练损失一路跌到0而验证损失反弹,基本就是过拟合了,朝上板斧去处理。
4.4 损失变成NaN怎么处理
训练着训练着,损失突然变成NaN,然后模型的预测结果全部乱掉。这个问题处理起来最需要耐心,但原因其实就那几类:
- 学习率太大:梯度更新迈的步子太大,参数飞出去了。把学习率从
0.001降到0.0001看是否解决。这是最高频的原因。 - 输入数据里混进NaN或无穷大:可以用
torch.isnan(data).any()检查输入数据。 - 梯度爆炸:参数瞬间变得极大,后面所有运算都溢出。可以加梯度裁剪缓解:
这行代码放在torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)loss.backward()之后、optimizer.step()之前,作用是限制梯度向量的长度,从根源上防止参数更新幅度过大。
4.5 常用调试习惯速查表
最后我把这两年积累的调试习惯整理成表格,你可以直接抄作业:
| 症状 | 首选排查动作 |
|---|---|
| forward报错维度不符 | 在forward里加print打印每步shape |
| 损失不下降 | 调大学习率0.01试一下,或者检查数据是否标准化 |
| 损失变成NaN | 降低学习率,检查输入是否有NaN,加梯度裁剪 |
| 训练和测试准确率差距大 | 模型处于过拟合,加Dropout或数据增强 |
| GPU利用率低 | 增大num_workers,检查是否有大量CPU预处理操作 |
| 显存溢出 | 降batch_size;检查是否忘了no_grad;torch.cuda.empty_cache() |
| CUDA不可用 | nvidia-smi看驱动;确认装的不是CPU版 |
| 每次测试结果忽高忽低 | 检查是否忘了调model.eval() |
聊到这儿,其实PyTorch最核心的思维模式就是“先跑通,再理解,再优化”。你不需要在第一次写代码时就完全掌握每一行背后的数学原理,但一定要在跑通之后,把那些训练循环里的“惯性动作”为什么存在搞清楚——为什么zero_grad、为什么no_grad、为什么要flatten。把这些“为什么”变成肌肉记忆之后,你再去看任何开源项目,会发现来来回回都是同一套流程:定义网络、加载数据、训练、评估、保存。到那时候,换个任务、换个模型结构,就只是改数据和改模块的事,真正核心的那套东西已经完全在你的脑子里了。