PyTorch入门实战：Tensor、自动求导与图像分类全流程

1. 先想明白你究竟在跟 PyTorch 学什么前阵子有个朋友问我说他想学 PyTorch但打开教程看到一堆nn.Module、DataLoader、backward()整个人就懵了。他说自己不是程序员只是做数据分析的被领导安排要训练一个图像分类模型问我应该怎么下手。我特别理解这种感觉。PyTorch 现在已经是深度学习领域使用最广的框架之一大量论文、开源项目、工业落地案例都跑在它上面。但大多数人第一次接触它问题往往不是PyTorch 怎么用而是我到底在学什么。如果你把 PyTorch 当成一个普通的 Python 库去学会越学越乱如果你把它当成一个自动求导机 张量计算器思路就会清晰很多。我习惯这样理解 PyTorch 的层级最底层是 Tensor负责存储数据和做各种数学运算类似 NumPy 数组的加强版但它能自动记录计算路径。第二层是autograd负责自动求导。深度学习的核心就是梯度下降而梯度就是损失函数对参数的导数手动推导非常痛苦PyTorch 帮你自动算了。第三层是torch.nn提供了神经网络组件比如卷积层、全连接层、损失函数。第四层是torch.optim封装了各类优化器比如 SGD、Adam 这样的参数更新算法。最上面才是DataLoader、训练循环、模型保存加载这些工程化内容。很多人一上来就盯着第五层写代码前四层靠复制粘贴。结果遇到问题就完全没法排查——因为你不知道数据流是怎么走的也不知道梯度是什么时候算出来的。这篇内容我打算从一个完整的图像分类项目切入带你从零跑通 PyTorch 的核心流程。中途我会重点解释为什么会这样设计也会把我在实际项目中踩过的坑直接摆出来。适合谁看呢两类人一类是刚接触深度学习、连 Tensor 和torch.Tensor都分不清的纯新手另一类是写了一段时间代码但还是靠拼凑的初学者希望通过一次完整梳理把整个体系打通。2. 装环境不劝退版本、显卡和那些缺一步就崩的细节很多人的 PyTorch 之旅不是倒在模型上而是倒在装环境上。说实话PyTorch 的安装已经比几年前友好太多了但网上教程五花八门新旧版本混杂照着一套老教程装很容易翻车。2.1 教你一个最稳的安装判断方法先别急着pip install torch先想清楚两个问题你用不用 GPU你的 CUDA 环境是什么情况打开终端按顺序看这几项python --version pip --version nvidia-smi重点说nvidia-smi这个命令。如果它提示找不到命令说明你机器上要么没有 NVIDIA 显卡要么驱动没装好。如果它正常输出了 GPU 信息那么看右上角的 CUDA Version,比如 12.1 或者 11.8这个数字代表你当前驱动支持的 CUDA 最高版本。然后去 PyTorch 官网找对应的安装命令。官网做的安装配置选择器已经非常直观了你只要把操作系统、包管理器、CUDA 版本选好它就会给出对应的pip命令。我见过最典型的坑是显卡驱动支持 CUDA 12.1但很多人选了 PyTorch 编译时用的 CUDA 11.8 版本结果也能跑只是某些算子可能会警告或者性能差一点。其实在大多数场景下用略低于驱动支持版本的 CUDA 轮子包问题不大因为 PyTorch 的 CUDA 支持是向后兼容的。但如果你装的是 CPU 版本然后发现torch.cuda.is_available()一直返回False就说明你没看清安装命令里的cpu标志。2.2 我推荐的环境管理方式作为过来人我非常建议你用 Python 虚拟环境来装 PyTorch不要让它在系统环境里裸奔。我见过太多人pip install torch之后某天装其他包时不小心把 torch 升级到不兼容版本整个项目直接跑不起来。我习惯用conda建环境conda create -n torch_env python3.10 conda activate torch_envPython 版本不用太激进3.9 或 3.10 都很稳。然后按照官网给出的pip命令装 torch。装完做个快速自检import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.device_count())如果这三行都正常环境就基本没问题。这里多说一句很多人看到torch.cuda.is_available()返回False就很慌,但其实如果你没有 NVIDIA 显卡这个False是正常的。CPU 也能跑 PyTorch只是慢一些。我第一年做深度学习的时候就是用 CPU 训练小型网络虽然慢但理解原理完全够用。还有一个很实际的建议装特定版本的 PyTorch 时尽量让 torch、torchvision、torchaudio 三个包保持同一个发布版本。比如你用2.1.0那这三个最好都升到2.1.0因为它们之间存在编译层面的对齐要求版本不一致经常会报出特别奇怪的错误。3. 拆掉 Tensor 和自动求导这堵墙在正式写模型前必须先打通 Tensor 和自动求导这层。这层不打通后面所有代码对你来说都是玄学——能跑但不知道为什么跑。3.1 Tensor 到底是什么用最简单的话说Tensor 是多维数组。标量是一个 0 维 Tensor向量是 1 维矩阵是 2 维图像数据是 4 维——形状一般是(批次大小, 通道数, 高, 宽)比如一批 32 张图像每张是 3 通道、224x224那 Tensor 的形状就是(32, 3, 224, 224)。你可以把它理解为带了很多额外功能的 NumPy 数组它能自动放到 GPU 上运算能记录操作历史从而自动求导。如果你会 NumPyTensor 上手很快import torch import numpy as np # 从列表创建 t1 torch.tensor([1.0, 2.0, 3.0]) # 从 NumPy 数组创建 arr np.array([1, 2, 3]) t2 torch.from_numpy(arr) # 创建全零矩阵 t3 torch.zeros(3, 4) print(t1.shape) # torch.Size([3]) print(t2.dtype) # torch.int64很多新手在数据处理阶段特别容易栽在 dtype 上。比如从 NumPy 读进来的是float64而 PyTorch 模型默认参数是float32一旦混用往往在计算损失时突然报错。解决办法就是显式转换t torch.from_numpy(arr).float()或者用torch.as_tensor(arr, dtypetorch.float32)。类似这种坑没有任何技术含量但如果你不知道根源会白白排查很久。3.2 自动求导为什么是 PyTorch 的立身之本机器学习训练的核心套路是前向传播算出预测值计算损失然后通过反向传播求出损失对每个参数的梯度最后用优化器更新参数。问题来了求梯度这件事对于简单函数你手动算算还行对于动辄几十层、上百层的神经网络手推是不可能的。而 PyTorch 的autograd机制会在你进行 Tensor 运算时自动构建一张计算图记住你是通过什么操作得到最终结果的。然后你调用.backward()时它会沿着这张图反方向把梯度算出来存到每个张量的.grad属性里。看一下最简单的例子x torch.tensor(2.0, requires_gradTrue) y x ** 2 y.backward() print(x.grad) # 输出 tensor(4.)这里y x * x导数在x 2处是 4。你看到的是 PyTorch 自动算出来了。你只需要在创建需要梯度变量的 Tensor 时标上requires_gradTrue或者用nn.Parameter包装。这里有个关键认知模型的权重参数是需要梯度的输入数据一般是不需要梯度的。所以如果你手写代码时给输入数据也设置了requires_gradTrue虽然不报错但计算图会变大内存和显存消耗上升训练变慢属于典型的不该做的动作。3.3 训练流程中梯度怎么一步步流动一个完整的训练循环我想先把全貌摆出来再拆解每一步在干什么import torch import torch.nn as nn import torch.optim as optim for epoch in range(num_epochs): for batch_data, batch_label in dataloader: batch_data batch_data.to(device) batch_label batch_label.to(device) # 前向传播 outputs model(batch_data) loss criterion(outputs, batch_label) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step()我第一次看到这段代码时非常不理解为什么每轮要zero_grad()。后来才明白PyTorch 的梯度是累积的如果你不清零上次算出来的梯度会和这次梯度叠加。这意味着如果不zero_grad()你的参数更新方向就是多次梯度的叠加不是一个正确的单步梯度。这是新手最容易忽略、但影响极大的细节。backward()是沿着计算图计算梯度。step()是让优化器拿着梯度去更新参数。这三种方法各干各的活很多人只会复制调用顺序却不知道去掉任何一个都会出问题。4. 用三小时跑通一个像样的图像分类 Demo光讲概念不讲实操就是耍流氓。这一章我带你手动搭一个用于图像分类的完整项目从数据准备到模型训练再到预测整个过程跟着敲一遍你会对 PyTorch 的工程链路有一个整体手感。4.1 数据怎么准备从原始图片到 DataLoader图像分类第一步是准备好数据。假设我们现在手里有一个猫狗分类的数据集目录长这样data/ train/ cat/ 1000 张图片 dog/ 1000 张图片 val/ cat/ 200 张图片 dog/ 200 张图片PyTorch 提供了一个非常常用的数据加载组合torchvision.datasets.ImageFolder负责按目录结构读取图片并打标签torchvision.transforms负责对图像做预处理torch.utils.data.DataLoader负责把数据打包成批次并自动打乱顺序。from torchvision import datasets, transforms from torch.utils.data import DataLoader transform transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rootdata/train, transformtransform) val_dataset datasets.ImageFolder(rootdata/val, transformtransform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4) print(train_dataset.classes) # [cat, dog]这里每个环节都值得细说。Resize((128, 128))把图片统一到固定尺寸因为模型要求固定尺寸输入。ToTensor()有两件事一是把PIL.Image或numpy.ndarray转成 Tensor二是把像素值从0~255缩放到0~1。Normalize用均值 0.485、标准差 0.229 这些数值这几个数不是随便拍的是 ImageNet 数据集的统计值。用它们在训练时能加快收敛属于行业通用做法。为什么我不在训练集里直接加随机裁剪、随机翻转这些数据增强因为在搭第一个 Demo 时数据增强会引入太多变量你很难判断模型性能变化到底来自哪里。先把最朴素的流程跑通后面再逐步加增强效果会比一步到位好得多。4.2 模型搭建别一上来就上 ResNet很多新手上来就想用 ResNet、EfficientNet 这种大网络我觉得没必要。理解 PyTorch 的首要目标是理解数据流和训练循环。所以我建议先手搭一个两层卷积的小网络既能看到nn.Conv2d的作用又不会因为网络太复杂导致排查困难。import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.conv1 nn.Conv2d(3, 16, kernel_size3, padding1) self.conv2 nn.Conv2d(16, 32, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) self.fc nn.Linear(32 * 16 * 16, num_classes) # 输入尺寸后续解释 def forward(self, x): x self.pool(F.relu(self.conv1(x))) x self.pool(F.relu(self.conv2(x))) x x.view(x.size(0), -1) x self.fc(x) return x输入是(32, 3, 128, 128)的批次。经过一次卷积加池化变成(32, 16, 64, 64)再经过一次变成(32, 32, 32, 32)。展平后每个样本是32 * 32 * 32 32768维所以全连接层输入就是 32768。我在上面代码里写了32 * 16 * 16那是我把输入尺寸改成 64x64 时的值实战里你最好用公式推一遍或者直接打印特征图形状。nn.Module是 PyTorch 模型的基础类。你的模型必须继承它才能真正融入自动求导体系。在__init__里定义子层在forward里定义前向计算逻辑。注意一个常见误解__init__里用nn.Conv2d这样的层因为它们是nn.Module的子类模型会自动管理它们的参数。如果你用普通 Python 列表存层参数不会被注册优化器也看不到它们。正确做法是用nn.Sequential或nn.ModuleList。4.3 训练模型的组装损失函数、优化器、训练循环网络定义好之后还需要选一个损失函数和一个优化器。对于图像分类这种离散标签任务最常用的是交叉熵损失nn.CrossEntropyLoss。这里有个特别容易搞错的设计这个损失函数本身已经包含了softmax步骤所以你不需要在模型的最后一层加 Softmax直接输出原始分数logits就行。model SimpleCNN(num_classes2) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)optimizer绑定的是model.parameters()也就是说优化器只知道模型里被注册的参数然后它会在每次step()里更新这些参数。lr0.001是 Adam 优化器一个非常常用的默认值大多数简单任务从 0.001 起步都比较稳妥。训练循环我之前已经写过这里我把关键部分再展开一下加上打印日志和验证逻辑def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total accuracy correct / total return avg_loss, accuracymodel.train()表示进入训练模式这个在简单网络里看起来没什么作用但之后用 Dropout 或 BatchNorm 就会出现训练和验证行为不一致的问题。torch.max(outputs, 1)是找出每一行最大值的索引也就是模型认为最可能的类别。.item()把 Tensor 转成 Python 数值方便打印。验证函数和训练函数几乎一样唯一不同的是要加with torch.no_grad():因为在验证阶段我们不需要计算梯度这样省显存也更快。还要调用model.eval()切到验证模式。4.4 跑起来看效果然后判断该往哪调如果上面代码都敲对了训练几个 epoch 后你会发现损失下降、准确率上升。我的经验是在猫狗这种简单二分类任务上用这个小网络训练大概十个 epoch验证准确率应该能到 90% 以上。如果发现损失不降或者准确率很低不要急着改网络结构按顺序检查几件事第一标签是否张冠李戴比如ImageFolder是按字母顺序分配类别的cat是 0、dog是 1确认没有反第二学习率是不是太大或太小太大损失会震荡太小收敛太慢第三数据预处理是不是出错了比如Normalize参数写反。跑通之后你可以试着把SimpleCNN换成torchvision.models.resnet18(weightstorchvision.models.ResNet18_Weights.DEFAULT)并修改最后全连接层输出来适配自己的类别数。迁移学习在这个体量的任务上效果会显著好于自己从零训练的小网络。这也是行业里最常见的做法——大多数人不会从零训练整个大模型而是用预训练模型做微调。5. Debug 才是主菜新手必掉的 5 个坑我在帮别人排查 PyTorch 代码的时候发现有些坑几乎每隔一段时间就会有人踩一次。下面这几类问题覆盖面很高每次遇到它们解决思路基本都一致。5.1 维度不匹配最常见的报错Expected shape ... but found ...这类错误的典型场景是把全连接层输入维度算错了。我把排查方法总结成一条在forward里临时打印每一层的输出形状。def forward(self, x): print(input:, x.shape) x self.pool(F.relu(self.conv1(x))) print(after conv1:, x.shape) x self.pool(F.relu(self.conv2(x))) print(after conv2:, x.shape) x x.view(x.size(0), -1) print(after flatten:, x.shape) x self.fc(x) return x调试完之后记得删掉这些print。依赖打印调试虽然看起来土但实话实说它比盯着权重初始化代码推演要快得多。我一般先用很小的批次跑一次确认数据流没问题再放开训练。5.2 梯度爆炸与梯度消失loss 变成 NaN训练着训练着 loss 突然变成nan这是很有冲击感的一类问题。最常见的原因有几个学习率太大导致梯度更新幅度过大数据里混入了NaN值网络初始化不良。最简单的排查方案是调低学习率比如从1e-3降到1e-4如果恢复正常说明大概率是学习率问题。另外可以在backward()后加一句torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这在训练 RNN、Transformer 类模型时几乎是标配操作可以防止梯度范数过大导致更新异常。5.3 CPU 和 GPU 张量混用PyTorch 的一个硬性要求是参与运算的张量必须在同一个设备上。最常见的报错信息是 Expected all tensors to be on the same device。原因很简单模型被.to(device)放到了 GPU但输入数据忘了.to(device)。这个问题的解决办法其实是在代码里统一管理设备。我的习惯是尽早定义device torch.device(cuda if torch.cuda.is_available() else cpu)然后模型和数据都在进入训练循环前调用.to(device)。5.4 训练集和验证集的模型模式搞混Dropout 和 BatchNorm 这两种层在训练和验证时行为不同。训练时 Dropout 会随机丢弃一部分神经元验证时不会BatchNorm 训练时使用当前批次的均值和方差验证时使用历史统计量。如果你在验证时忘了model.eval()验证结果会不稳定时高时低。我见过不少人训练完直接在训练脚本里顺手算了一个验证准确率结果比自己预期低不少可能是 Dropout 引入了随机性然后就开始怀疑模型有问题。其实只要在验证前面加一行model.eval()就解决了。5.5 过拟合训练准确率高但验证准确率低训练集准确率升到 100%验证集只有 70%这种状态叫过拟合意思是模型死记硬背了训练样本没有学到通用的规律。最简单的应对方式有三板斧加数据增强比如随机水平翻转、随机裁剪、色彩抖动加 Dropout或者改用预训练模型做迁移学习。对于图像分类任务在这三者里面优先尝试数据增强成本最低、效果最明显。这些小问题单独来看都不难但它们组合起来的排查成本很高。如果你刚开始写 PyTorch我可以给你一个非常实用的定位思路如果把模型换成极小的网络问题依然存在那说明问题多半不在模型结构而在数据处理或训练流程。6. 从能跑到跑得好我的几个私人习惯最后聊一些不常写进教程但实际很有用的经验。我在调试所有深度学习项目时都会先做一个单批次过拟合测试只拿一个批次的数据训练几十步看模型能不能把这一批数据完全记住损失降到极低。能做到这一步说明你的代码链路是通的做不到说明哪里有问题。这一步还能帮你快速发现损失是否下降、参数是否更新是最推荐做的 sanity check。然后是模型保存。PyTorch 保存模型通常有两种做法torch.save(model.state_dict(), model.pth)只保存参数torch.save(model, model.pth)保存整个模型。我强烈建议用第一种因为它只保存张量数据体积小跨设备、跨版本兼容性更好。加载时先重新定义模型结构再load_state_dict。torch.save(model.state_dict(), model.pth) model SimpleCNN(num_classes2) model.load_state_dict(torch.load(model.pth, map_locationdevice)) model.eval()map_locationdevice这个参数很关键。你在 GPU 上训练的模型在 CPU 机器上加载时必须指定map_locationcpu不然会报设备不匹配的错误。还有关于随机种子的设置。如果你希望实验可复现最好固定随机种子import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)不过要说明的是在 GPU 上做并行运算时同样的种子也不能保证 100% 结果一致因为 CUDA 的某些操作本身有随机性。所以不要追求绝对的复现固定种子只是帮你把大多数变量控制住。回看整个 PyTorch 的学习路径你会发现它其实没有想象中那么难。核心只围绕一件事张量计算、自动求导、参数更新。把这些环节拆开理解再通过一个完整的小项目串起来比起背 API 要高效得多。我现在做一个新任务时也依然会按最朴素的流程先跑通再逐步优化。这个习惯救过我非常多次。