PyTorch 深度学习入门:从张量、自动求导到 CNN 实战与模型部署
2026/9/9 0:12:15 网站建设 项目流程

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_env

3.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 matplotlib

5.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.9906

5.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单精度浮点数823约 1e-38 到 3e38默认精度,通用稳定训练初始权重、CPU 推理、对精度敏感的层
fp16半精度浮点数510约 6e-5 到 65504范围小,小数值容易下溢,大数值容易上溢混合精度训练、GPU 推理加速
bf16脑浮点数87与 fp32 相同的大范围范围大,但尾数位数少,精度更低大模型训练、梯度累积、分布式训练
tf32Tensor Float 32810与 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()为 FalsePyTorch 版本和 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 的过程,看到过准确率曲线从抖动到平稳的轨迹,才算真正入门。建议把这篇文章里的两个项目代码保存成自己的模板,后续遇到新的任务,只需要替换数据加载和模型结构,训练流程基本不需要大改。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询