PyTorch 深度学习入门到底该怎么学?这个问题几乎每隔几天就会在技术社群里出现一次。很多初学者从“安装 PyTorch”开始,到跑通 MNIST 结束,中间隔着一整座山:张量、自动求导、Dataset、DataLoader、训练循环、验证指标……每一步单独看都能看懂,合在一起就不知道怎么串起来。
这篇文章的目标,就是帮你把这座山拆成一条可以照着走的路。我会从 PyTorch 最核心的几个设计讲起,把它和传统编程思维的差异说清楚,然后用两个完整的实战项目(一个线性回归、一个手写数字识别 CNN),把“数据处理 -> 模型搭建 -> 训练 -> 验证 -> 保存”这一整套流程跑通。最后,我还会结合深度学习模型部署中最常见的 fp32、fp16、bf16、tf32 浮点数格式,讲清楚什么时候该用哪种精度,帮你提前避坑。
读完这篇文章,你不需要再到处搜“PyTorch 环境怎么搭建”“CNN 的 shape 怎么算”“训练不收敛怎么办”这类零散问题。你只需要照着一路敲下来,就能拥有自己的第一个可用于真实任务的训练脚本。
1. PyTorch 到底解决什么问题
先说判断:PyTorch 是目前学术界和工业界交接最顺畅的深度学习框架。它真正降低的,是“想法 -> 实验 -> 验证”这个循环的成本。
在 PyTorch 出现之前,用 TensorFlow 1.x 写一个模型,需要先把计算图静态定义好,再通过 Session 去执行。这种“先构图、后执行”的方式,对刚接触深度学习的人来说非常不直观:你想打印中间某个张量的值,都要按“占位符 -> 运算节点 -> feed_dict”的流程走一遍,调试一个 bug 可能要花掉半天。
PyTorch 把这件事彻底改掉了。它的核心设计是“动态计算图”:每一次前向传播,都会即时构建一张计算图,所有中间结果就是你日常使用的普通张量。你可以像写普通 Python 代码一样打断点、打印、修改变量,梯度信息会随着计算图自动保留和更新。用一句被用滥了但很准确的话说:PyTorch 让深度学习代码写起来更像 Python,而不是像一个被框架绑架的 DSL。
这个设计直接影响了整个生态。研究人员需要快速改模型结构、试新想法,PyTorch 的动态图让这种改动的成本变得极低。工业界需要把训练好的模型部署到服务器,PyTorch 后来又推出了 TorchScript、torch.compile、TorchServe 等工具链,把研究和生产的距离拉近了。这就是为什么从 2018 年之后,顶会论文里 PyTorch 的比例一路走高,很多公司内部的模型训练基础设施也开始以 PyTorch 为底座。
但是,这不代表 PyTorch 没有学习成本。它的难点不在于 API 数量,而在于思维方式:张量的 shape 变化、自动求导的机制、模块化设计理念、训练循环的构成。这些不是靠背 API 能解决的,需要靠理解加动手。接下来我从最底层的张量讲起。
2. PyTorch 核心概念与工作原理
2.1 张量(Tensor):深度学习时的“数据容器”
深度学习中,所有数据在进入模型之前,都会被表示为张量。可以把张量理解成 NumPy 的多维数组,但它有两个关键区别:一是支持 GPU 加速计算,二是参与自动求导。
从维度上理解张量最直观:
- 0 维张量是一个标量(scalar),比如损失值 loss;
- 1 维张量是一个向量(vector),比如一句话经过词嵌入后的 256 维向量;
- 2 维张量是一个矩阵(matrix),比如一批 64 条样本、每条样本 128 个特征,shape 就是 (64, 128);
- 3 维张量可以表示时序或图像单通道数据,比如一批 64 张 28x28 的灰度图,shape 是 (64, 1, 28, 28) 或者 (64, 28, 28);
- 4 维张量最常用的是彩色图像批数据,比如一批 64 张 224x224 的 RGB 图,shape 是 (64, 3, 224, 224)。
在实际代码中,创建张量非常简单:
import torch # 从 Python 列表创建 x = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) print(x.shape) # torch.Size([2, 2]) print(x.dtype) # torch.float32(默认浮点类型) # 创建全 0 / 全 1 / 随机张量 zeros = torch.zeros(2, 3) ones = torch.ones(2, 3) randn = torch.randn(2, 3) # 标准正态分布随机数 # 张量运算与 NumPy 非常相似 y = x + ones z = torch.matmul(x, x.T) # 矩阵乘法,x.T 表示转置在深度学习模型里,最常见的一个新手困惑是:什么时候该用 view / reshape / permute?这里先记一个结论:view 和 reshape 是改变张量“形状”但不改变“元素排列顺序”的操作,permute 是交换维度。如果你在处理图像数据,想把形状从 (batch, height, width, channel) 改成 (batch, channel, height, width),就要用 permute,而不是 view,因为 view 不会移动数据,会直接把元素按顺序切分,导致通道错乱。
# 假设 input 的 shape 是 (batch=4, height=32, width=32, channel=3) # 这是常见的 HWC 图像布局 input_hwc = torch.randn(4, 32, 32, 3) # 错误做法:用 view 硬改 # wrong = input_hwc.view(4, 3, 32, 32) # 正确做法:用 permute 交换维度顺序 input_chw = input_hwc.permute(0, 3, 1, 2) print(input_chw.shape) # torch.Size([4, 3, 32, 32])2.2 自动求导:PyTorch 的“发动机”
深度学习训练的本质,是反复执行“前向传播计算损失 -> 反向传播计算梯度 -> 用梯度更新模型参数”。自动求导(autograd)是 PyTorch 之所以好用的核心机制。
PyTorch 的做法是:任何一个张量,只要你设置requires_grad=True,所有基于它产生的运算都会被记录到一张计算图中。当你调用loss.backward()时,PyTorch 会从 loss 出发,沿计算图反向传播,给每个设置了requires_grad=True的张量填充.grad属性。
import torch # 定义一个需要梯度的参数 w = torch.tensor([2.0], requires_grad=True) x = torch.tensor([3.0]) # 前向计算 y = w * x # y = 6 loss = (y - 1) ** 2 # loss = (6-1)^2 = 25 # 反向传播 loss.backward() # 查看梯度:d(loss)/d(w) = 2 * (y - 1) * x = 2 * 5 * 3 = 30 print(w.grad) # tensor([30.]) # 更新参数后,必须清空梯度 w.grad.zero_()实操里,我们不会手动管理每个参数的梯度,而是把可训练参数封装进torch.nn.Module里。nn.Module不仅管理参数,还提供parameters()方法,方便优化器统一更新。PyTorch 官方教程里常说的“训练三件套”,就是:loss.backward()计算梯度、optimizer.step()更新参数、optimizer.zero_grad()清空上一轮梯度。这三个步骤顺序不能乱。
很多人会问:为什么每次更新前要先zero_grad()?因为 PyTorch 的梯度是累加的(accumulate),如果你不清空,下一轮的梯度会叠加到上一轮上,等于每一步都在用“历史总梯度”更新参数,模型永远不会收敛。
2.3 Dataset 与 DataLoader:数据从“乱七八糟的文件”变成“模型能吃的批数据”
在真实项目中,数据通常不在内存里,而是散落在磁盘上的图片、文本、CSV 文件。PyTorch 提供了两个数据处理核心类:
torch.utils.data.Dataset:定义“如何从原始数据中取出一个样本”。torch.utils.data.DataLoader:负责把多个样本组织成 batch、打乱顺序、并行加载。
一个典型的自定义 Dataset 长这样:
import torch from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, data, labels): self.data = data self.labels = labels def __len__(self): return len(self.data) def __getitem__(self, idx): # 返回一个样本(x, y) return self.data[idx], self.labels[idx] # 模拟数据 data = torch.randn(1000, 32) labels = torch.randint(0, 10, (1000,)) dataset = MyDataset(data, labels) dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2) for batch_x, batch_y in dataloader: print(batch_x.shape) # torch.Size([32, 32]) print(batch_y.shape) # torch.Size([32]) break这里的shuffle=True在训练时非常重要,它能让每个 epoch 的样本顺序不同,避免模型学到样本顺序相关的伪规律。num_workers控制多进程加载数据的数量,训练数据量大时能显著减少 CPU 等待时间,但新手阶段设置成 0 或 1 就可以。
2.4 nn.Module:所有模型的“积木”
nn.Module是 PyTorch 神经网络模型的基类。你搭建的任何模型,都应该继承它。它自动完成参数注册、parameters()收集、.to(device)迁移等工作,还能通过嵌套子模块组装出复杂的网络结构。
import torch.nn as nn class MyMLP(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super().__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x model = MyMLP(32, 64, 10) print(model) print([p.shape for p in model.parameters()])forward方法定义了数据从输入到输出的计算过程。PyTorch 会在调用model(x)时自动调用forward。在forward里,你可以像写普通函数一样写条件判断、循环、打印,这就是动态图带来的灵活性。
3. PyTorch 环境准备与安装
3.1 环境规划
在动手安装之前,先想清楚三件事:操作系统、Python 版本、是否使用 GPU。
- 操作系统:Windows、Linux、macOS 都可以,但生产环境绝大多数情况下是 Linux。Ubuntu 22.04 / 24.04 是深度学习领域最常见的系统版本。
- Python 版本:PyTorch 新版对 Python 版本有要求,一般建议使用 Python 3.9 - 3.12。不建议用系统自带 Python,强烈推荐 Anaconda 或 Miniconda 做环境隔离。
- GPU:如果你的电脑有 NVIDIA 显卡,安装 CUDA 版本的 PyTorch 能极大加速训练。如果没有 GPU,可以先安装 CPU 版本跑通逻辑,但深度学习模型的训练速度会慢很多。
3.2 使用 conda 创建独立环境
深度学习项目之间经常有依赖冲突。比如项目 A 依赖 PyTorch 2.0,项目 B 却需要 PyTorch 1.13,如果共用环境,很容易出现“装完 B,A 跑不了了”的窘境。独立环境是解决这个问题最有效的方式。
# 创建名为 pytorch_env 的 conda 环境,指定 Python 版本 conda create -n pytorch_env python=3.10 -y # 激活环境 conda activate pytorch_env3.3 安装 PyTorch
打开 PyTorch 官网(pytorch.org),找到 Get Started 页面,它会根据你的操作系统、包管理工具、CUDA 版本生成对应的安装命令。尽量不要手动输入版本号,直接复制官方生成的命令最靠谱。
常见命令示例:
# CPU 版本 pip install torch torchvision torchaudio # CUDA 12.1 版本(请以官网生成命令为准) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果你在国内网络环境下下载较慢,可以配置 PyTorch 的国内镜像源来加速。这里以清华源为例:
pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple如果你的显卡是 AMD 或 Intel,也可以使用 ROCm 或其他后端的安装方式。但需要注意,多数深度学习生态默认优化的是 NVIDIA CUDA,AMD 用户需要额外确认自己的 PyTorch 版本支持 ROCm。从材料中可以看到“AMD 安装 PyTorch CUDA”这类搜索词存在,说明这也是新手常见疑惑点。如果是这种情况,建议去 PyTorch 官网查看对应平台的安装说明,不要盲目套用别人的命令。
3.4 验证安装是否成功
安装完成后,一定不要急着写代码,先验证环境是否正常。创建一个test_env.py文件,输入以下内容:
import torch import torchvision print("PyTorch 版本:", torch.__version__) print("torchvision 版本:", torchvision.__version__) # 检查是否有可用的 CUDA GPU print("CUDA 是否可用:", torch.cuda.is_available()) if torch.cuda.is_available(): print("GPU 型号:", torch.cuda.get_device_name(0)) print("GPU 显存:", torch.cuda.get_device_properties(0).total_memory / 1024**3, "GB")运行:
python test_env.py如果torch.cuda.is_available()返回 True,说明 GPU 环境正常。如果是 False,不要慌,先看一下 PyTorch 版本和 CUDA 驱动是否匹配。一个常见坑是:显卡驱动版本足够新,但 PyTorch 安装的 CUDA 运行时版本和驱动不完全匹配,这时需要更新驱动或更换 PyTorch 的 CUDA 版本。
4. PyTorch 完整实战一:线性回归
线性回归是深度学习的“Hello World”。它虽然简单,但能让你完整看到训练流程的每一个环节。
4.1 项目目标
我们用 PyTorch 拟合一个一元二次函数y = 2*x^2 + 1加入少量噪声后的数据。模型是一个只有一个隐藏层的 MLP,配合 ReLU 激活函数,就能拟合非线性关系。
4.2 完整代码
import torch import torch.nn as nn import torch.optim as optim import matplotlib.pyplot as plt # 1. 生成模拟数据 torch.manual_seed(42) x = torch.linspace(-3, 3, 200).reshape(-1, 1) y = 2 * x**2 + 1 + 0.1 * torch.randn_like(x) # 2. 定义模型 class MLP(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 1) ) def forward(self, x): return self.net(x) model = MLP() # 3. 定义损失函数和优化器 criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.01) # 4. 训练循环 epochs = 500 loss_history = [] for epoch in range(epochs): optimizer.zero_grad() output = model(x) loss = criterion(output, y) loss.backward() optimizer.step() loss_history.append(loss.item()) if (epoch + 1) % 100 == 0: print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.6f}") # 5. 预测并可视化 model.eval() with torch.no_grad(): y_pred = model(x) plt.figure(figsize=(8, 5)) plt.scatter(x.numpy(), y.numpy(), alpha=0.5, label="真实数据") plt.plot(x.numpy(), y_pred.numpy(), color="red", linewidth=2, label="模型预测") plt.legend() plt.title("PyTorch 线性回归拟合结果") plt.show()4.3 关键逻辑解释
torch.manual_seed(42):固定随机种子,保证每次运行结果一致,便于调试。nn.Sequential:把多个层按顺序组合,适合这种线性堆叠的网络结构。criterion = nn.MSELoss():回归任务最常用的损失函数,计算预测值和真实值的均方误差。optimizer = optim.Adam(model.parameters(), lr=0.01):Adam 优化器,是目前最常用的默认选择。lr是学习率,控制每次参数更新的步长。model.eval()和with torch.no_grad():在验证/推理阶段必须使用。它告诉 PyTorch 不要更新 BatchNorm 和 Dropout 的统计状态,同时关闭梯度计算,节省内存。
这个例子虽然简单,但五脏俱全。如果你能看懂它的每一行,理解训练循环的四个步骤,你就已经掌握了 PyTorch 训练所有模型的标准骨架。
5. PyTorch 完整实战二:手写数字识别 CNN
线性回归只是热身,真正的深度学习入门项目必然是卷积神经网络(CNN)。这里用 MNIST 数据集做一个端到端的完整项目。
5.1 为什么用 MNIST
MNIST 是深度学习的“果蝇”:数据规模小、下载容易、任务标准明确。一个简单的 CNN 就能达到 99% 以上的准确率,非常适合新手用来验证自己的代码流程是否正确,而不会被训练时间和调参分散注意力。
5.2 环境准备
pip install torch torchvision matplotlib5.3 完整代码
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import matplotlib.pyplot as plt # 1. 定义数据预处理 transform = transforms.Compose([ transforms.ToTensor(), # 将 PIL 图像转为张量,并归一化到 [0, 1] transforms.Normalize((0.1307,), (0.3081,)) # 用 MNIST 数据集的均值和标准差做标准化 ]) # 2. 下载并加载数据集 train_dataset = datasets.MNIST( root="./data", train=True, transform=transform, download=True ) test_dataset = datasets.MNIST( root="./data", train=False, transform=transform, download=True ) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False) # 3. 定义 CNN 模型 class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv_layers = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2), # 28x28 -> 14x14 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2) # 14x14 -> 7x7 ) self.fc_layers = nn.Sequential( nn.Flatten(), # 64 * 7 * 7 -> 3136 nn.Linear(64 * 7 * 7, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, 10) ) def forward(self, x): x = self.conv_layers(x) x = self.fc_layers(x) return x model = SimpleCNN() # 4. 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 5. 定义训练函数 def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in dataloader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) _, predicted = torch.max(outputs, 1) correct += (predicted == labels).sum().item() total += labels.size(0) avg_loss = total_loss / total accuracy = correct / total return avg_loss, accuracy # 6. 定义测试函数 def evaluate(model, dataloader, criterion, device): model.eval() total_loss, correct, total = 0.0, 0, 0 with torch.no_grad(): for images, labels in dataloader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() * images.size(0) _, predicted = torch.max(outputs, 1) correct += (predicted == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total # 7. 训练并验证 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) print("使用设备:", device) epochs = 5 for epoch in range(1, epochs + 1): train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device) test_loss, test_acc = evaluate(model, test_loader, criterion, device) print(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f} | " f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.4f}") # 8. 保存模型 torch.save(model.state_dict(), "mnist_cnn.pth") print("模型已保存到 mnist_cnn.pth")运行结果大致是这样的(不同随机种子和硬件上会有小幅波动):
使用设备: cuda Epoch 1: Train Loss=0.2459, Train Acc=0.9266 | Test Loss=0.0727, Test Acc=0.9781 Epoch 2: Train Loss=0.0791, Train Acc=0.9756 | Test Loss=0.0472, Test Acc=0.9849 Epoch 3: Train Loss=0.0586, Train Acc=0.9819 | Test Loss=0.0362, Test Acc=0.9881 Epoch 4: Train Loss=0.0484, Train Acc=0.9848 | Test Loss=0.0323, Test Acc=0.9890 Epoch 5: Train Loss=0.0416, Train Acc=0.9871 | Test Loss=0.0285, Test Acc=0.99065.4 模型的 shape 变化
很多人第一次写 CNN 时,最头疼的是全连接层输入的维度怎么算。这里给出这个模型的 shape 变化:
- 输入图像:(64, 1, 28, 28)
- 第一个 Conv + ReLU + MaxPool:(64, 32, 14, 14)
- 第二个 Conv + ReLU + MaxPool:(64, 64, 7, 7)
- Flatten 后:(64, 64 * 7 * 7) = (64, 3136)
- 第一个全连接层输出:(64, 128)
- Dropout 后:(64, 128)
- 第二个全连接层输出:(64, 10)
要特别注意,padding=1配合kernel_size=3时,卷积不会改变特征图的宽高。MaxPool2d(2, 2)会把宽高各减半。图中数据的 shape 变换是 CNN 学习中最基础的内容,建议新手在纸上把这个流程推导一遍。
5.5 训练效果验证
训练完成后,你可以随机取几张测试集中的图片,看看模型的预测结果。这是判断模型是否真的学到了东西的直观方式:
import matplotlib.pyplot as plt # 加载测试集中的一个批次 data_iter = iter(test_loader) images, labels = next(data_iter) # 取前 6 张图 sample_images = images[:6] sample_labels = labels[:6] # 用训练好的模型预测 model.eval() with torch.no_grad(): outputs = model(sample_images) _, predicted = torch.max(outputs, 1) # 可视化 fig, axes = plt.subplots(2, 3, figsize=(8, 6)) for idx, ax in enumerate(axes.ravel()): img = sample_images[idx].squeeze().numpy() ax.imshow(img, cmap="gray") ax.set_title(f"真实: {sample_labels[idx].item()} | 预测: {predicted[idx].item()}") ax.axis("off") plt.tight_layout() plt.show()如果预测结果和真实标签一致,说明模型已经基本掌握手写数字识别的规律。
6. 深度学习模型部署必知:fp32、fp16、bf16、tf32 怎么选
很多人在本地训练时没有感知,但一旦进入模型部署和性能优化阶段,就会被浮点数精度问题卡住。这里有几个名词,是模型部署和面试中的高频考点:fp32、fp16、bf16、tf32。简单来说,它们决定了模型在 GPU 上存储权重、中间激活值和计算梯度时使用的数值精度和范围。
6.1 四种浮点数格式对比
| 格式 | 全称 | 指数位 | 尾数位 | 数值范围 | 精度特点 | 典型场景 |
|---|---|---|---|---|---|---|
| fp32 | 单精度浮点数 | 8 | 23 | 约 1e-38 到 3e38 | 默认精度,通用稳定 | 训练初始权重、CPU 推理、对精度敏感的层 |
| fp16 | 半精度浮点数 | 5 | 10 | 约 6e-5 到 65504 | 范围小,小数值容易下溢,大数值容易上溢 | 混合精度训练、GPU 推理加速 |
| bf16 | 脑浮点数 | 8 | 7 | 与 fp32 相同的大范围 | 范围大,但尾数位数少,精度更低 | 大模型训练、梯度累积、分布式训练 |
| tf32 | Tensor Float 32 | 8 | 10 | 与 fp32 基本相同 | 介于 fp32 和 fp16 之间 | NVIDIA Ampere 架构 GPU 上的加速计算 |
6.2 为什么不能无脑用 fp16
fp16 的数值范围大约是 65504,很多深度学习训练过程中产生的中间结果很容易超出这个范围,产生上溢出或下溢出。解决这个问题的标准方案是混合精度训练(Mixed Precision Training):权重副本保持 fp32,前向和反向计算用 fp16,Loss Scale 将损失值放大后再缩回。PyTorch 自带的torch.cuda.amp模块已经把这一套封装好了。
从实际部署角度看,fp16 的推理速度通常比 fp32 快,内存占用减半,但如果没有处理好精度损失,模型输出的质量可能会明显下降。比较稳妥的做法是:先用 fp32 跑出基线准确率,再切换到 fp16/bf16/tf32 对比效果,确认没有严重精度回退后再上线。
6.3 新手选型建议
- 训练阶段:优先用 fp32,代码简单,调试方便。如果你的模型较大、训练很慢,再考虑自动混合精度。
- 部署阶段:NVIDIA T4、A10、A100 等 GPU 上,fp16 和 int8 量化是主流加速方式。
- 大模型训练:bf16 因为保留了与 fp32 相同的指数范围,成为训练 10B 以上参数模型时的首选。
- 对精度敏感的层(如 BatchNorm、损失函数计算),建议保留 fp32,避免数值问题。
7. PyTorch 常见问题与排查思路
新手在跑深度学习项目时,会遇到的问题其实高度重复。这里把我们最常见的几个问题整理成表格,并补充一些判断思路。
| 问题现象 | 可能原因 | 排查方式 | 解决方案 |
|---|---|---|---|
ModuleNotFoundError: No module named 'torch' | 没有安装 PyTorch,或安装到了其他 conda 环境 | 在终端输入python -c "import torch"确认 | conda activate激活正确环境,重新安装 PyTorch |
CUDA out of memory | 某个进程占用了大量显存,或 batch_size 过大 | 使用nvidia-smi查看显存占用;检查是否有残留进程 | 减小 batch_size,降低图片分辨率,或kill残留 Python 进程 |
RuntimeError: shape '[64, 10]' is invalid for input of size ... | 全连接层输入维度计算错误 | 打印模型每层的输出 shape,逐步定位 | 根据公式手算 shape,或使用torchsummary.summary(model, (1, 28, 28))查看 |
| 训练 Loss 是 NaN | 学习率过大、数据里有 NaN、或者模型除零 | 打印 loss 和模型输出的数值范围 | 降低学习率,检查数据预处理是否产生了 NaN |
| 训练 Loss 不下降 | 学习率过小、模型结构错误、优化器参数没设置对 | 用少量样本过拟合测试 | 先在小数据上跑通,验证模型能否过拟合,再换大图 |
GPU 可用但torch.cuda.is_available()为 False | PyTorch 版本和 CUDA 驱动不匹配 | 查看nvidia-smi里的 CUDA 版本和 PyTorch 的版本 | 到 PyTorch 官网选择匹配的 CUDA 版本重新安装 |
| 同一个代码在别人机器上能跑,自己机器报错 | 环境不一致:Python 版本、依赖版本、CUDA 版本 | 对比pip list和 Python 版本 | 建议使用 conda 环境和 requirements.txt 锁定版本 |
补充一个非常重要的排错方法:让模型先在小数据集上过拟合。当你能在一个 batch 的数据上把训练准确率跑到 100%,说明模型结构和优化器设置是通的。如果连一个 batch 都拟合不了,问题大概率出在模型定义或数据 pipeline,而不是数据量不足。
8. PyTorch 最佳实践与工程建议
8.1 代码结构分层
不要把所有代码都塞进一个文件。到项目后期,你会非常感谢当初把代码分成这几个模块的自己:
project/ ├── config.py # 全局超参数配置 ├── data_loader.py # Dataset 和数据预处理 ├── model.py # 模型定义 ├── train.py # 训练脚本 ├── evaluate.py # 验证 / 测试脚本 ├── utils.py # 通用工具函数 └── requirements.txt # 依赖版本列表这种分层方式的核心价值是“职责单一”。模型结构、数据逻辑、训练流程各自独立,改动任何一个模块都不影响其他模块,也方便多个人协作。
8.2 保存和加载模型的坑
很多人只记住一句“用 torch.save 保存模型”,但保存方式选错了,后面加载时会踩大坑。
# 推荐方式:只保存模型参数(state_dict) torch.save(model.state_dict(), "model.pth") # 加载模型时,必须先重新实例化模型结构 model = SimpleCNN() model.load_state_dict(torch.load("model.pth")) model.eval()不要直接保存整个模型对象:
# 不推荐:保存整个模型对象,容易因为代码版本变化导致加载失败 # torch.save(model, "model.pth")只保存state_dict的好处是,模型结构由代码控制,文件只保存权重数据,跨版本和跨环境兼容性更好。加载后记得调用model.eval(),把 Dropout 和 BatchNorm 切到推理模式,否则预测结果会和训练模式不一致。
8.3 随机种子管理
深度学习结果的可复现性一直被忽视,直到你发现“昨天还能跑通,今天换台机器结果完全不同”。解决办法是在训练脚本开头固定所有随机种子:
import random import numpy as np import torch def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False set_seed(42)注意,设置cudnn.deterministic = True会略微降低训练速度,但能保证同一份代码在同样的硬件和环境下输出一致。对于需要写论文或做实验对比的场景,这是必须的。
8.4 日志与监控
训练深度模型是一个长时间过程,如果看不到中间状态,你根本不知道模型是在收敛还是在振荡。建议在训练循环里至少打印以下信息:当前 epoch、训练 loss、验证 loss、训练准确率、验证准确率、学习率、已用时间。有条件的话,用 TensorBoard 记录这些指标,效果更直观。
from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter("runs/experiment_1") # 在每个 epoch 结束后 writer.add_scalar("Loss/train", train_loss, epoch) writer.add_scalar("Loss/test", test_loss, epoch) writer.add_scalar("Accuracy/test", test_acc, epoch)8.5 生产环境注意事项
如果是为生产环境训练模型,有几点容易被忽略:
- 训练环境和部署环境分离。训练用高性能 GPU 服务器,推理用专门的推理服务,不要混用。
- 数据集版本管理。训练集的数据如果发生了更新,要能追溯到具体版本,否则模型迭代时无法对比效果。
- 模型版本管理。每个训练出来的模型文件,最好用一个唯一标识(比如训练时间、git commit、数据集版本)命名。
- 业务上线前,一定要在真实分布的数据上校验模型效果,而不是只依赖测试集指标。
9. 总结与后续学习方向
这篇文章从 PyTorch 的基础概念出发,讲清楚了张量、自动求导、Dataset、DataLoader、nn.Module 这几个核心组成部分,并用线性回归和 MNIST CNN 两个完整项目演示了标准的训练流程。同时,我们把 fp32、fp16、bf16、tf32 这几种浮点数格式做了对比,分析了它们在训练和部署中的不同选择场景。对新手来说,这篇文章最大的价值是帮你建立了一个“最小可行知识体系”:先跑通一个完整项目,再在真实项目中逐步补充细节。
下一步的学习路径,我建议按这个顺序推进:
- 先复现本文的 MNIST 项目,在不看代码的情况下自己写一遍;
- 把 MNIST 换成 CIFAR-10,体验多类别彩色图像分类和更大的数据量;
- 学习如何使用 torchvision.models 里预训练好的 ResNet、VGG 做迁移学习;
- 掌握
torch.utils.tensorboard和可视化工具,建立模型训练的“仪表盘”意识; - 开始阅读经典模型的源码,从
torchvision.models.resnet出发,理解残差结构、BatchNorm 的实现细节; - 当训练足够成熟后,关注模型部署:ONNX 导出、TensorRT 推理加速、量化。
学习 PyTorch 最忌讳的是“只看不敲”。你可以在脑海里推导几千行代码,但只有真正在终端里跑过一次模型训练,经历过 loss 从 2.3 降到 0.01 的过程,看到过准确率曲线从抖动到平稳的轨迹,才算真正入门。建议把这篇文章里的两个项目代码保存成自己的模板,后续遇到新的任务,只需要替换数据加载和模型结构,训练流程基本不需要大改。