从零到一:基于PyTorch与VGG16的猫狗分类实战指南
1. 环境配置GPU与CPU的选择刚开始接触深度学习时最让我纠结的就是到底用GPU还是CPU。说实话我以前一直觉得GPU配置特别复杂直到自己动手试了一次才发现其实就跟装普通软件差不多。这里分享下我的真实踩坑经历帮你少走弯路。如果你有NVIDIA显卡强烈建议用GPU版本。我用的是一块GTX 1060显卡实测训练速度比CPU快8倍不止。关键就三步先装CUDA再装对应版本的PyTorch最后装cuDNN。记得一定要检查版本匹配我第一次就栽在这里装了CUDA 11.0却配了PyTorch 1.7.0需要CUDA 10.2结果死活跑不起来。具体操作步骤到NVIDIA官网下载CUDA Toolkit推荐10.2或11.3安装时记得勾选添加到PATH在PyTorch官网用官方命令安装比如conda install pytorch torchvision cudatoolkit10.2 -c pytorch验证安装在Python里跑torch.cuda.is_available()返回True就成功了没有显卡也不用慌。CPU版本安装更简单直接pip install torch torchvision就行。虽然训练慢点但小数据集完全够用。我之前用i7-9750H跑2000张图的猫狗分类一个epoch大概3分钟完全可以接受。2. 数据准备从Kaggle到自定义数据集网上教程动不动就用Kaggle的25000张图对新手太不友好了。我自己建了个精简版数据集只要400张图200猫200狗效果也不错。这里教你怎么快速构建自己的数据集。首先去Kaggle下载原始数据集搜索dogs-vs-cats解压后你会看到train文件夹里有一大堆命名为cat.0.jpg、dog.1.jpg的图片。我写了个自动分拣脚本import os import shutil def organize_dataset(original_dir, target_dir, sample_size200): os.makedirs(f{target_dir}/train/cats, exist_okTrue) os.makedirs(f{target_dir}/train/dogs, exist_okTrue) os.makedirs(f{target_dir}/test/cats, exist_okTrue) os.makedirs(f{target_dir}/test/dogs, exist_okTrue) # 复制训练集 for i in range(sample_size): shutil.copy(f{original_dir}/cat.{i}.jpg, f{target_dir}/train/cats) shutil.copy(f{original_dir}/dog.{i}.jpg, f{target_dir}/train/dogs) # 复制测试集用300-400的图片 for i in range(300, 300 sample_size//2): shutil.copy(f{original_dir}/cat.{i}.jpg, f{target_dir}/test/cats) shutil.copy(f{original_dir}/dog.{i}.jpg, f{target_dir}/test/dogs)重点来了数据增强这是提升模型泛化能力的关键。我用的是PyTorch的transforms组合from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomRotation(20), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])3. 模型构建迁移学习实战技巧直接从头训练CNN简直是灾难我用VGG16的预训练模型做迁移学习准确率轻松突破95%。这里有几个教科书不会告诉你的细节冻结卷积层只训练最后的全连接层可以防止小数据集上的过拟合model models.vgg16(pretrainedTrue) for param in model.parameters(): # 冻结所有参数 param.requires_grad False自定义分类头原模型是1000类分类我们要改成二分类model.classifier nn.Sequential( nn.Linear(25088, 256), nn.ReLU(), nn.Dropout(0.5), # 对抗过拟合 nn.Linear(256, 2) )渐进式解冻先训练全连接层再解冻部分卷积层# 训练5轮后解冻最后两个卷积块 for name, param in model.features.named_parameters(): if 24 in name or 26 in name: param.requires_grad True训练时我用的是带momentum的SGD优化器比Adam更适合迁移学习optimizer optim.SGD([ {params: model.classifier.parameters(), lr: 1e-3}, {params: model.features.parameters(), lr: 1e-4} ], momentum0.9)4. 训练与评估避坑指南第一次训练时我的准确率卡在60%上不去后来发现是学习率设太高。这里分享我的调参经验学习率从1e-4开始尝试用ReduceLROnPlateau动态调整Batch SizeGPU显存8G的话设16-32比较安全Early Stopping连续3轮验证集loss不下降就停止我写的训练循环包含这些关键功能from torch.optim.lr_scheduler import ReduceLROnPlateau scheduler ReduceLROnPlateau(optimizer, min, patience2) best_acc 0 for epoch in range(20): # 训练阶段 model.train() for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() # 验证阶段 model.eval() with torch.no_grad(): correct 0 total 0 for inputs, labels in val_loader: outputs model(inputs) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total if acc best_acc: torch.save(model.state_dict(), best_model.pth) best_acc acc scheduler.step(loss)可视化特别重要用Matplotlib画损失曲线和准确率曲线plt.figure(figsize(12,4)) plt.subplot(121) plt.plot(train_losses, labeltrain) plt.plot(val_losses, labelval) plt.title(Loss curve) plt.legend() plt.subplot(122) plt.plot(train_acc, labeltrain) plt.plot(val_acc, labelval) plt.title(Accuracy curve) plt.legend()5. 模型部署实用预测技巧训练好的模型怎么用我开发了个带界面的预测脚本支持单张图片和批量预测from PIL import Image import matplotlib.pyplot as plt def predict_image(img_path, model, transform): img Image.open(img_path).convert(RGB) img_t transform(img).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(img_t) _, pred torch.max(outputs, 1) # 可视化显示 plt.imshow(img) plt.title(fPredicted: {cat if pred0 else dog}) plt.axis(off) plt.show() return pred.item()遇到模糊难辨的图片怎么办我加了个置信度显示probs torch.nn.functional.softmax(outputs, dim1)[0] print(fCat: {probs[0]:.2%}, Dog: {probs[1]:.2%})最后提醒几个容易翻车的点预测前务必做和训练时相同的预处理记得调用model.eval()关闭dropout和batchnorm用with torch.no_grad()禁用梯度计算节省内存