ResNet18二分类实战：云端GPU 10分钟训练宠物识别模型-平芜编程栈

ResNet18二分类实战：云端GPU 10分钟训练宠物识别模型

引言

开宠物店的朋友最近遇到了一个头疼的问题：店里每天要处理大量猫狗照片，手动分类不同品种耗时费力。找外包公司报价动辄上万元，作为小本生意实在难以承受。其实用AI技术完全可以自己解决这个问题，今天我就带大家用ResNet18模型，在云端GPU上快速训练一个宠物识别分类器。

ResNet18是深度学习领域经典的图像分类模型，就像给电脑装上了一双能自动识别猫狗品种的"智能眼睛"。即使你完全没有AI经验也不用担心，跟着本文操作，从零开始到完成训练只需10分钟。我们会使用CSDN星图镜像广场提供的预置环境，省去复杂的配置过程，直接进入模型训练环节。

1. 环境准备：5分钟搞定AI开发环境

1.1 选择适合的云端GPU环境

训练AI模型就像炒菜需要好锅一样，GPU就是我们的"智能炒锅"。普通电脑的CPU处理图像数据太慢，而GPU能大幅加速训练过程。在CSDN星图镜像广场，我们可以直接选择预装了PyTorch和CUDA的镜像，省去自己配置环境的麻烦。

推荐选择以下配置： - 镜像类型：PyTorch 1.12 + CUDA 11.6 - GPU型号：至少4GB显存（如NVIDIA T4） - 存储空间：20GB以上

1.2 快速启动云端环境

登录CSDN星图平台后，按照以下步骤操作：

在镜像广场搜索"PyTorch"
选择官方提供的PyTorch基础镜像
点击"立即部署"按钮
等待1-2分钟环境初始化完成

部署成功后，你会获得一个类似Jupyter Notebook的在线开发环境，所有必要的软件都已预装好。

2. 数据准备：整理你的宠物照片

2.1 收集猫狗品种照片

好的数据是AI模型的基础，就像教小朋友认动物需要准备清晰的图片一样。我们需要两类数据：

猫的品种照片（建议至少200张）
狗的品种照片（建议至少200张）

可以从以下几个渠道获取： - 自己宠物店日常拍摄的照片 - 公开数据集（如Kaggle上的宠物数据集） - 网络搜索（注意版权）

2.2 整理数据文件夹结构

在云端环境中创建一个项目文件夹，按如下结构组织：

pet_classification/ ├── train/ │ ├── cat/ │ │ ├── cat1.jpg │ │ ├── cat2.jpg │ │ └── ... │ └── dog/ │ ├── dog1.jpg │ ├── dog2.jpg │ └── ... └── val/ ├── cat/ └── dog/

train文件夹用于训练，约占全部数据的80%
val文件夹用于验证，约占20%

3. 模型训练：10分钟打造专属分类器

3.1 加载预训练模型

ResNet18已经在百万级ImageNet数据集上预训练过，我们可以直接利用这些学到的特征，就像用已经认识很多动物的老师来专门学习识别猫狗。

import torch import torchvision.models as models # 加载预训练的ResNet18模型 model = models.resnet18(pretrained=True) # 修改最后一层全连接层，适应我们的二分类任务 num_features = model.fc.in_features model.fc = torch.nn.Linear(num_features, 2) # 2代表猫狗两个类别

3.2 准备数据加载器

PyTorch提供了方便的工具来加载和预处理图像数据：

from torchvision import datasets, transforms # 定义数据增强和归一化 data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), # ResNet18标准输入尺寸 transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 加载数据集 image_datasets = { 'train': datasets.ImageFolder('pet_classification/train', data_transforms['train']), 'val': datasets.ImageFolder('pet_classification/val', data_transforms['val']) } # 创建数据加载器 dataloaders = { 'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=32, shuffle=True), 'val': torch.utils.data.DataLoader(image_datasets['val'], batch_size=32, shuffle=False) }

3.3 配置训练参数

训练AI模型就像教小朋友学习，需要设置合适的学习节奏：

import torch.optim as optim # 定义损失函数和优化器 criterion = torch.nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 学习率调度器 scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

3.4 开始训练模型

现在可以启动训练过程了，GPU会让这个步骤非常快速：

# 将模型转移到GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device) # 训练5个epoch（完整遍历数据集5次） num_epochs = 5 for epoch in range(num_epochs): # 训练阶段 model.train() for inputs, labels in dataloaders['train']: inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 验证阶段 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in dataloaders['val']: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Epoch {epoch+1}/{num_epochs}, 准确率: {100 * correct / total:.2f}%') scheduler.step()

4. 模型使用与优化技巧

4.1 保存训练好的模型

训练完成后，我们可以把模型保存下来供后续使用：

torch.save(model.state_dict(), 'pet_classifier.pth')

4.2 加载模型进行预测

使用时可以这样加载模型并进行预测：

# 加载保存的模型 model.load_state_dict(torch.load('pet_classifier.pth')) model.eval() # 单张图片预测函数 def predict_image(image_path): image = Image.open(image_path) image = data_transforms['val'](image).unsqueeze(0) image = image.to(device) with torch.no_grad(): outputs = model(image) _, predicted = torch.max(outputs, 1) return 'cat' if predicted.item() == 0 else 'dog'