用PyTorch从零搭建CNN:我的第一个猫狗分类器实战(附完整代码与数据集)
2026/5/31 2:32:39 网站建设 项目流程

用PyTorch从零搭建CNN:我的第一个猫狗分类器实战(附完整代码与数据集)

第一次接触深度学习时,我被那些能自动识别图片内容的算法深深吸引。作为一个刚入门的新手,最让我困惑的是如何将教科书上的理论知识转化为实际可运行的代码。直到亲手用PyTorch完成这个猫狗分类项目,才真正理解了CNN(卷积神经网络)的运作机制。本文将分享这个过程中积累的实战经验,包括那些教科书不会告诉你的"坑"和解决方案。

1. 项目准备与环境搭建

在开始编码之前,我们需要确保开发环境配置正确。推荐使用Anaconda创建独立的Python环境,避免依赖冲突。以下是关键软件版本要求:

  • Python 3.8+
  • PyTorch 1.12+
  • torchvision 0.13+
  • OpenCV 4.5+
# 创建conda环境 conda create -n pytorch_cnn python=3.8 conda activate pytorch_cnn # 安装PyTorch(根据CUDA版本选择) pip install torch torchvision torchaudio

提示:如果使用GPU加速训练,需提前安装对应版本的CUDA和cuDNN。可以通过nvidia-smi命令查看显卡驱动版本。

数据集准备是第一个容易踩坑的环节。Kaggle的"Dogs vs Cats"数据集包含25,000张图片,但实际训练时我们可能只需要部分数据。我整理了一个精简版数据集(约4000张图片),下载后目录结构应如下:

data/ ├── train/ │ ├── cat/ │ └── dog/ └── test/ ├── cat/ └── dog/

2. 数据预处理与加载

原始图像尺寸不一,直接输入网络会导致问题。我们需要统一处理为224x224像素,并进行标准化。PyTorch的torchvision.transforms模块提供了便捷的预处理方法:

from torchvision import transforms # 定义训练集和测试集的转换流程 train_transform = transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), # 数据增强:随机裁剪 transforms.RandomHorizontalFlip(), # 数据增强:水平翻转 transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

自定义Dataset类时,我遇到了路径处理的典型问题。以下是优化后的实现:

from torch.utils.data import Dataset import os class CatDogDataset(Dataset): def __init__(self, root_dir, transform=None): self.root_dir = root_dir self.transform = transform self.classes = ['cat', 'dog'] self.images = [] # 遍历目录收集样本路径 for label, class_name in enumerate(self.classes): class_dir = os.path.join(root_dir, class_name) for img_name in os.listdir(class_dir): self.images.append((os.path.join(class_dir, img_name), label)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path, label = self.images[idx] image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) return image, label

3. CNN模型设计与实现

经过多次实验,我设计了一个适合初学者的简化版CNN结构。相比复杂模型,这个版本更容易理解且训练速度快:

import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.features = nn.Sequential( # 卷积层1:输入3通道,输出16通道,3x3卷积核 nn.Conv2d(3, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2), # 卷积层2:输入16通道,输出32通道 nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2), # 卷积层3:输入32通道,输出64通道 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2) ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(64*28*28, 512), # 注意计算特征图尺寸 nn.ReLU(), nn.Dropout(0.5), # 防止过拟合 nn.Linear(512, 2) ) def forward(self, x): x = self.features(x) x = self.classifier(x) return x

注意:全连接层的输入尺寸需要根据前面的卷积层输出计算。一个常见错误是忽略这个计算,导致运行时维度不匹配。

4. 模型训练与优化

训练过程中有几个关键参数需要仔细调整:

参数推荐值说明
学习率0.001-0.0001太大导致震荡,太小收敛慢
Batch Size32-64根据GPU内存调整
Epochs10-20观察验证集准确率变化
import torch.optim as optim from tqdm import tqdm def train_model(model, criterion, optimizer, num_epochs=10): best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch+1}/{num_epochs}') print('-' * 10) # 每个epoch有训练和验证阶段 for phase in ['train', 'val']: if phase == 'train': model.train() # 训练模式 else: model.eval() # 评估模式 running_loss = 0.0 running_corrects = 0 # 迭代数据 for inputs, labels in tqdm(dataloaders[phase]): inputs = inputs.to(device) labels = labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 反向传播+优化仅在训练阶段 if phase == 'train': loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # 深度拷贝模型 if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) print(f'Best val Acc: {best_acc:4f}') model.load_state_dict(best_model_wts) return model

5. 模型评估与可视化

训练完成后,我们需要评估模型在测试集上的表现。除了准确率,混淆矩阵能提供更多信息:

from sklearn.metrics import confusion_matrix import seaborn as sns def evaluate_model(model, test_loader): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for inputs, labels in test_loader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 计算混淆矩阵 cm = confusion_matrix(all_labels, all_preds) plt.figure(figsize=(8,6)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['Cat', 'Dog'], yticklabels=['Cat', 'Dog']) plt.xlabel('Predicted') plt.ylabel('True') plt.title('Confusion Matrix') plt.show() return cm

可视化卷积层的特征图有助于理解CNN的工作原理:

def visualize_feature_maps(model, image_tensor): # 获取各卷积层的输出 activations = [] def hook_fn(module, input, output): activations.append(output) hooks = [] for layer in [model.features[0], model.features[3], model.features[6]]: hooks.append(layer.register_forward_hook(hook_fn)) # 前向传播 model.eval() with torch.no_grad(): _ = model(image_tensor.unsqueeze(0).to(device)) # 移除hook for hook in hooks: hook.remove() # 可视化特征图 fig, axes = plt.subplots(3, 8, figsize=(20,8)) for i, activation in enumerate(activations): act = activation.squeeze().cpu().numpy() for j in range(min(8, act.shape[0])): axes[i,j].imshow(act[j], cmap='viridis') axes[i,j].axis('off') plt.tight_layout() plt.show()

6. 常见问题与解决方案

在实际项目中,我遇到了以下典型问题及解决方法:

问题1:GPU内存不足

现象:训练时出现CUDA out of memory错误

解决方案

  • 减小batch size(如从64降到32)
  • 使用混合精度训练:
    from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() with autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

问题2:模型过拟合

现象:训练准确率高但验证准确率低

解决方案

  • 增加数据增强(随机旋转、颜色抖动等)
  • 添加Dropout层
  • 使用早停法(Early Stopping)
  • 尝试更简单的模型结构

问题3:训练loss不下降

现象:多个epoch后loss仍无明显变化

解决方案

  • 检查学习率是否合适
  • 验证数据预处理是否正确
  • 确认模型结构是否有误(如激活函数缺失)
  • 检查梯度更新是否正常:
    # 在训练循环中添加 for name, param in model.named_parameters(): if param.grad is None: print(f"No gradient for {name}") else: print(f"{name} grad norm: {param.grad.norm().item()}")

7. 项目扩展与优化方向

完成基础版本后,可以考虑以下优化:

  1. 更高效的模型结构

    • 尝试ResNet、EfficientNet等现代架构
    • 使用预训练模型进行迁移学习
  2. 高级训练技巧

    # 学习率预热 from torch.optim.lr_scheduler import LambdaLR warmup_epochs = 5 scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: (epoch+1)/warmup_epochs if epoch < warmup_epochs else 1)
  3. 部署应用

    • 使用Flask创建Web接口
    • 转换为ONNX格式优化推理速度
    • 使用TorchScript进行序列化

完整项目代码和数据集已整理在GitHub仓库中,包含更多详细注释和实用工具函数。通过这个项目,我深刻体会到实践是学习深度学习的最佳方式——那些看似复杂的理论,在代码实现过程中会变得清晰而直观。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询