PyTorch通用开发实战案例：微调ResNet全流程部署指南

1. 引言：为什么选择这个环境做ResNet微调？

你是不是也经历过这样的场景：每次开始一个新项目，都要花半天时间配环境、装依赖、解决版本冲突？尤其是用PyTorch做图像分类任务时，光是把CUDA、cuDNN和各种库对齐就让人头大。

今天我们要做的，是在一个开箱即用的PyTorch通用开发环境中，完成一次完整的ResNet模型微调实战。这个环境基于官方PyTorch镜像构建，预装了Pandas、Numpy、Matplotlib和JupyterLab等常用工具，系统干净无冗余缓存，还配置了阿里云和清华源，下载依赖飞快。

更重要的是，它支持CUDA 11.8 / 12.1，完美适配RTX 30/40系列显卡以及A800/H800等企业级GPU。这意味着你不需要再为驱动兼容问题发愁，拿到手就能直接跑训练。

我们以经典的ResNet模型为例，带你从数据准备、模型加载、微调训练到最终部署，走完深度学习项目的完整流程。无论你是刚入门的新手，还是想快速验证想法的开发者，这套方案都能帮你省下大量时间。

2. 环境准备与快速验证

2.1 启动开发环境

假设你已经通过容器或虚拟机方式启动了PyTorch-2.x-Universal-Dev-v1.0镜像，进入终端后第一件事就是确认GPU是否正常挂载。

运行以下命令：

nvidia-smi

你应该能看到类似下面的输出：

+-----------------------------------------------------------------------------+ | NVIDIA-SMI 535.129.03 Driver Version: 535.129.03 CUDA Version: 12.2 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 NVIDIA RTX 4090 Off | 00000000:01:00.0 Off | Off | | 30% 45C P8 10W / 450W | 0MiB / 24576MiB | 0% Default | +-------------------------------+----------------------+----------------------+

这说明你的GPU已经被正确识别。

接下来验证PyTorch能否使用CUDA：

python -c "import torch; print(f'PyTorch版本: {torch.__version__}'); print(f'GPU可用: {torch.cuda.is_available()}'); print(f'当前设备: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'No GPU'}')"

理想输出如下：

PyTorch版本: 2.3.0 GPU可用: True 当前设备: NVIDIA RTX 4090

如果看到这些信息，恭喜你，环境已经ready！

3. 数据准备：构建自定义图像分类数据集

3.1 数据结构设计

我们要微调ResNet来做一个简单的图像分类任务——区分猫、狗和兔子三类动物。你可以用自己的图片，也可以从公开数据集中抽取样本（比如Kaggle上的Pet Dataset）。

数据目录结构建议如下：

data/ ├── train/ │ ├── cat/ │ ├── dog/ │ └── rabbit/ └── val/ ├── cat/ ├── dog/ └── rabbit/

每类训练集放100~200张图即可，验证集放20~30张。图片尽量统一尺寸（如224x224），格式为.jpg或.png。

3.2 使用TorchVision进行数据增强与加载

PyTorch提供了强大的torchvision模块，我们可以用它轻松实现数据预处理。

创建一个data_loader.py文件：

from torchvision import datasets, transforms from torch.utils.data import DataLoader # 定义训练和验证的数据变换 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 加载数据集 train_dataset = datasets.ImageFolder('data/train', transform=train_transform) val_dataset = datasets.ImageFolder('data/val', transform=val_transform) # 创建DataLoader train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4) print(f"训练样本数: {len(train_dataset)}") print(f"验证样本数: {len(val_dataset)}") print(f"类别名称: {train_dataset.classes}")

这段代码做了几件事：

对训练集做了随机裁剪、水平翻转等增强操作
验证集只做中心裁剪，保证评估一致性
使用标准ImageNet均值和方差进行归一化
批量加载数据，并开启多线程加速

运行后你会看到：

训练样本数: 450 验证样本数: 75 类别名称: ['cat', 'dog', 'rabbit']

4. 模型搭建：加载预训练ResNet并修改分类头

4.1 选择合适的ResNet版本

PyTorch内置了多个ResNet变体：ResNet18、ResNet34、ResNet50、ResNet101等。对于小数据集，推荐使用ResNet18或ResNet34，避免过拟合。

我们这里选用ResNet18：

import torch import torch.nn as nn from torchvision.models import resnet18, ResNet18_Weights # 获取预训练权重（推荐使用Weights API） weights = ResNet18_Weights.DEFAULT model = resnet18(weights=weights) # 冻结所有卷积层参数 for param in model.parameters(): param.requires_grad = False # 修改最后的全连接层，适配我们的3分类任务 num_features = model.fc.in_features model.fc = nn.Linear(num_features, 3) # 3个类别 # 将模型移到GPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) print(model.fc) # 查看新的分类层

输出应为：

Linear(in_features=512, out_features=3, bias=True)

提示：ResNet18_Weights.DEFAULT会自动下载最高质量的预训练权重（通常是ImageNet-1K训练过的）。这是PyTorch 2.0+推荐的做法，比旧版的pretrained=True更灵活。

4.2 只训练最后的分类层

由于我们冻结了主干网络的所有参数，只有最后一层fc需要更新。这样可以大幅减少训练时间和显存占用。

检查一下哪些参数需要梯度更新：

for name, param in model.named_parameters(): if param.requires_grad: print(name)

你应该只会看到：

fc.weight fc.bias

完美！说明其他层都已成功冻结。

5. 训练流程：定义损失函数、优化器与训练循环

5.1 设置训练超参数

import torch.optim as optim from torch.optim.lr_scheduler import StepLR # 超参数 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.fc.parameters(), lr=1e-3) scheduler = StepLR(optimizer, step_size=5, gamma=0.1) # 每5轮学习率×0.1 num_epochs = 10

这里我们只让fc层的参数参与优化，所以传入的是model.fc.parameters()。

5.2 编写训练与验证循环

from tqdm import tqdm best_acc = 0.0 for epoch in range(num_epochs): print(f"\nEpoch {epoch+1}/{num_epochs}") print('-' * 30) # 训练阶段 model.train() running_loss = 0.0 corrects = 0 for inputs, labels in tqdm(train_loader, desc="Training"): inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) _, preds = torch.max(outputs, 1) corrects += (preds == labels).sum().item() epoch_loss = running_loss / len(train_dataset) epoch_acc = corrects / len(train_dataset) print(f"Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}") # 验证阶段 model.eval() val_loss = 0.0 val_corrects = 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) val_loss += loss.item() * inputs.size(0) _, preds = torch.max(outputs, 1) val_corrects += (preds == labels).sum().item() val_loss = val_loss / len(val_dataset) val_acc = val_corrects / len(val_dataset) print(f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}") # 保存最佳模型 if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), "best_resnet18_pet.pth") print(f" 模型已保存，准确率: {best_acc:.4f}") scheduler.step() print(f"\n 最佳验证准确率: {best_acc:.4f}")

训练过程中你会看到类似这样的输出：

Epoch 1/10 ------------------------------ Training: 100%|██████████| 15/15 [00:08<00:00, 1.76it/s] Train Loss: 0.8921 Acc: 0.6222 Val Loss: 0.4123 Acc: 0.8400 模型已保存，准确率: 0.8400

经过10轮训练，通常能达到85%以上的准确率。

6. 模型推理：如何在新图片上做预测

训练完成后，我们来测试一下模型的实际效果。

创建一个predict.py文件：

from PIL import Image import torch from torchvision import transforms # 加载模型 model = resnet18(weights=None) model.fc = nn.Linear(512, 3) model.load_state_dict(torch.load("best_resnet18_pet.pth")) model = model.to(device) model.eval() # 图像预处理 transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 读取测试图片 img_path = "test_dog.jpg" img = Image.open(img_path).convert("RGB") img_tensor = transform(img).unsqueeze(0).to(device) # 增加batch维度 # 推理 with torch.no_grad(): output = model(img_tensor) _, pred = torch.max(output, 1) classes = ['cat', 'dog', 'rabbit'] print(f"预测结果: {classes[pred.item()]}")

运行后输出：