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 cudatoolkit=10.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_size=200): os.makedirs(f"{target_dir}/train/cats", exist_ok=True) os.makedirs(f"{target_dir}/train/dogs", exist_ok=True) os.makedirs(f"{target_dir}/test/cats", exist_ok=True) os.makedirs(f"{target_dir}/test/dogs", exist_ok=True) # 复制训练集 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(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])3. 模型构建:迁移学习实战技巧
直接从头训练CNN简直是灾难,我用VGG16的预训练模型做迁移学习,准确率轻松突破95%。这里有几个教科书不会告诉你的细节:
- 冻结卷积层:只训练最后的全连接层,可以防止小数据集上的过拟合
model = models.vgg16(pretrained=True) 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} ], momentum=0.9)4. 训练与评估:避坑指南
第一次训练时我的准确率卡在60%上不去,后来发现是学习率设太高。这里分享我的调参经验:
- 学习率:从1e-4开始尝试,用ReduceLROnPlateau动态调整
- Batch Size:GPU显存8G的话设16-32比较安全
- Early Stopping:连续3轮验证集loss不下降就停止
我写的训练循环包含这些关键功能:
from torch.optim.lr_scheduler import ReduceLROnPlateau scheduler = ReduceLROnPlateau(optimizer, 'min', patience=2) 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, label='train') plt.plot(val_losses, label='val') plt.title('Loss curve') plt.legend() plt.subplot(122) plt.plot(train_acc, label='train') plt.plot(val_acc, label='val') 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(f'Predicted: {"cat" if pred==0 else "dog"}') plt.axis('off') plt.show() return pred.item()遇到模糊难辨的图片怎么办?我加了个置信度显示:
probs = torch.nn.functional.softmax(outputs, dim=1)[0] print(f'Cat: {probs[0]:.2%}, Dog: {probs[1]:.2%}')最后提醒几个容易翻车的点:
- 预测前务必做和训练时相同的预处理
- 记得调用model.eval()关闭dropout和batchnorm
- 用with torch.no_grad()禁用梯度计算节省内存