从零开始：PyTorch让我爱上了深度学习

初识PyTorch，为何选择它？

我第一次接触 PyTorch 是在两年前。当时我们团队刚接手一个图像分类项目，原本是用 TensorFlow 做的，但在模型调优过程中遇到了一些“难搞”的问题，特别是在调试和可视化方面总觉得束手束脚。

于是，我在网上看到了越来越多人推荐 PyTorch，说是调试灵活、上手简单、社区活跃。抱着试试看的心态，我在周末抽了点时间跑了个 MNIST 分类 demo。结果只花了两个小时，就跑出了比之前 TensorFlow 模型还好的准确率，那一刻我就被震撼到了。

从此之后，我和 PyTorch 的“恋爱故事”正式开始了。

---

业务背景：我们需要一个高效的图像分类系统

项目背景

我们是一家做工业质检的公司，客户需要我们开发一个视觉缺陷检测系统，用来判断产品表面是否存在划痕、裂纹或异物等瑕疵。这类场景中，传统的基于规则的算法已经无法满足复杂多变的纹理特征识别需求，所以我们决定引入深度学习技术。

最初版本中，我们使用的是 ResNet18，并在私有数据集上进行微调。训练时采用交叉熵损失函数，并对输入图片进行了简单的增强（旋转、翻转等）。整个系统的部署流程大致如下：

1. 在 GPU 上训练模型
2. 使用 torchscript 导出为 .pt 文件
3. 部署到边缘设备上推理

听上去挺顺利？其实中间踩了不止一个坑。

---

遇到的第一个大坑：训练不收敛！

问题描述

当我们第一次训练模型的时候，无论是调整学习率还是更换优化器，loss 总是在某个值上下波动，始终不能很好地收敛。

我一开始怀疑是数据集划分的问题，后来又检查模型是否加载正确、权重是否冻结得当……一通折腾后还是没解决问题。

更让人头疼的是，TensorBoard 显示 loss 波动剧烈，但 acc 并没有明显提升，模型像一直在“原地打转”。

解决思路

这个时候我突然想到：是不是数据处理环节出了问题？

我重新审视了数据预处理的部分。原来我们在 DataLoader 中用了 num_workers > 0，而在 Windows 系统下，默认的 multiprocessing 方式会 fork 子进程加载数据，而某些库在这种情况下会有问题（尤其是 OpenCV），造成部分 batch 的 label 被错误标记。

解决办法：

• 将 num_workers=0，这样就不会用多进程处理数据
• 或者将 DataLoader 的 num_workers 设置为合适的数值，并确保所有自定义 Dataset 类中的 transform 操作都是线程安全的

改完后，loss 开始正常下降，accuracy 也逐步上升，训练过程终于走上正轨。

---

代码实践：PyTorch 入门的核心步骤

为了帮助大家更好地理解，下面是一个简化版的图像分类训练代码片段。这个例子假设你已经准备好了自己的数据集，并组织成了以下结构：

dataset/
├── train/
│   ├── class1/
│   ├── class2/
├── val/
│   ├── class1/
│   ├── class2/

数据准备与加载

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

train_dataset = datasets.ImageFolder(root='dataset/train', transform=transform)
val_dataset = datasets.ImageFolder(root='dataset/val', transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

模型构建（ResNet18）

import torch
import torchvision.models as models

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = models.resnet18(pretrained=True)
model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 修改输出层为2个类别
model.to(device)

损失函数与优化器

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

训练循环简例

for epoch in range(5):  # 只训5个epoch为例
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

---

更多实战经验：模型训练的那些“细节控”

1. 数据增强的重要性

刚开始我们没做太多数据增强，模型很容易过拟合。后来加上了随机裁剪、颜色扰动、仿射变换等方式，验证集准确率提升了将近 5%。

transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2),
transforms.RandomAffine(degrees=10, translate=(0.1, 0.1)),

这些看似“鸡肋”的小操作，在实际中往往效果显著，尤其在样本量不够的情况下。

---

2. Learning Rate 调整策略

初期我们固定使用 0.001 的学习率，但后面发现模型很快就陷入局部最优。于是尝试加入学习率调度器：

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2)

或者使用余弦退火：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

这对避免模型卡在 plateau 区域很有帮助。

---

3. 模型保存与恢复

训练中断是常态，特别是跑大型模型时。保存 checkpoint 是基本操作：

torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, "checkpoint.pth")

恢复训练也很简单：

checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

---

那些年踩过的雷

雷区一：tensor device 不一致

最常见的是你在 CPU 上构建 tensor，却试图用 GPU 进行运算，这时候会出现：

Expected object of device type cuda but got device type cpu for argument

解决方案很简单：统一放到同一个 device 上：

inputs = inputs.to(device)
labels = labels.to(device)

不过要注意：如果数据本身已经在 GPU 上，重复 .to(device) 也不会报错，只是会浪费性能。

---

雷区二：模型推理时没加 model.eval()

这会导致 Dropout 或 BatchNorm 出现训练阶段的行为，严重影响预测结果的一致性。

---

雷区三：导出模型忘记 trace or script

如果你要部署 PyTorch 模型到生产环境，一定要注意导出方式。例如：

script_model = torch.jit.script(model)
torch.jit.save(script_model, "traced_model.pt")

或者使用 trace：

dummy_input = torch.rand(1, 3, 224, 224).to(device)
traced_model = torch.jit.trace(model, dummy_input)

别忘了输入尺寸必须符合预期。

---

结果与收益：准确率 + 部署效率双双提升

经过一个多月的迭代和调优，我们的系统最终达到了 97.3% 的验证集准确率。相比最初的 89%，这是一个质的飞跃。

更重要的是，通过 PyTorch 提供的灵活性，我们可以快速实现模型实验、热更新、模型对比等功能。这也让我们团队在后续几个新项目中都优先选择了 PyTorch。

现在我们已经有了一套完整的训练框架模板，能复用在多个任务上：比如目标检测、图像分割、甚至文本相关任务也能兼容。效率直接起飞。

---

给新手的一些建议

1. 不要怕写 bug，PyTorch 最擅长的就是 debug。动态图机制让你每一行都能看到中间结果，非常适合研究型工作。

2. 善用 Jupyter Notebook 快速验证想法。你可以快速写一段代码，看看 Tensor 输出长啥样，模型跑起来会不会 crash。

3. 坚持用 logging 替代 print。尤其是在多人协作时，清晰的日志记录非常重要。

4. 不要追求“一步到位”。很多新手上来就想跑 ImageNet 级别的网络，但其实先从小任务练起更容易建立信心。建议从 CIFAR-10、Fashion-MNIST 开始练手。

5. 学会读源码。PyTorch 官方文档虽然丰富，但有时候还是要靠翻源码来理解函数行为。

6. 关注官方博客和技术论坛。PyTorch 社区更新很快，新工具层出不穷，掌握最新特性会让你事半功倍。

---

写在最后：PyTorch 是我见过最容易爱上深度学习的框架

以前我总觉得 AI 开发很神秘、门槛很高，直到遇见了 PyTorch。

它给了我一种“我能掌控一切”的感觉。我可以随时 inspect 模型的中间变量，可以在训练中临时打断程序分析异常，也可以一边写代码一边实时看到效果。

这种交互性和透明度，是静态图框架难以企及的优势。

如果你正在学习深度学习，或者想换一个更顺手的框架，不妨给 PyTorch 一次机会。也许就像我一样，它也会成为你最爱的那个“工具伴侣”。