logo

开发者热搜

从零开始：使用PyTorch完成Kaggle猫狗图像识别全流程解析

作者：十万个为什么2025.09.18 17:44浏览量：26

简介：本文详细介绍如何使用PyTorch框架完成Kaggle猫狗图像识别任务，涵盖数据加载、模型构建、训练优化及预测部署全流程，适合初学者及进阶开发者参考。

数据准备与预处理

数据集获取与结构分析

Kaggle猫狗分类数据集包含25,000张训练图片（12,500猫/12,500狗）和12,500张测试图片。数据按目录组织为train/cat、train/dog和test三个子目录。建议使用Kaggle API下载数据集，或通过kaggle competitions download -c dogs-vs-cats命令获取。

数据增强策略

为提升模型泛化能力，需实现以下数据增强：

  1. 随机水平翻转（概率0.5）
  2. 随机旋转（-15°~+15°）
  3. 随机缩放裁剪（224x224区域）
  4. 颜色抖动（亮度/对比度/饱和度调整）
  5. 标准化（均值[0.485,0.456,0.406]，标准差[0.229,0.224,0.225]）

PyTorch中可通过torchvision.transforms.Compose实现：

  1. train_transform = transforms.Compose([
  2.     transforms.RandomResizedCrop(224),
  3.     transforms.RandomHorizontalFlip(),
  4.     transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
  5.     transforms.ToTensor(),
  6.     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  7. ])

自定义数据加载器

使用torch.utils.data.Dataset创建自定义数据集类：

  1. class CatDogDataset(Dataset):
  2.     def __init__(self, root_dir, transform=None):
  3.         self.root_dir = root_dir
  4.         self.transform = transform
  5.         self.classes = ['cat', 'dog']
  6.         self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
  7.         self.images = []
  9.         for cls in self.classes:
  10.             cls_dir = os.path.join(root_dir, cls)
  11.             for img_name in os.listdir(cls_dir):
  12.                 self.images.append((os.path.join(cls_dir, img_name), self.class_to_idx[cls]))
  14.     def __len__(self):
  15.         return len(self.images)
  17.     def __getitem__(self, idx):
  18.         img_path, label = self.images[idx]
  19.         image = Image.open(img_path).convert('RGB')
  21.         if self.transform:
  22.             image = self.transform(image)
  24.         return image, label

模型架构设计

基础CNN实现

对于初学者，可构建包含4个卷积块的简单CNN：

  1. class SimpleCNN(nn.Module):
  2.     def __init__(self):
  3.         super(SimpleCNN, self).__init__()
  4.         self.features = nn.Sequential(
  5.             nn.Conv2d(3, 32, kernel_size=3, padding=1),
  6.             nn.ReLU(inplace=True),
  7.             nn.MaxPool2d(kernel_size=2, stride=2),
  8.             nn.Conv2d(32, 64, kernel_size=3, padding=1),
  9.             nn.ReLU(inplace=True),
  10.             nn.MaxPool2d(kernel_size=2, stride=2),
  11.             nn.Conv2d(64, 128, kernel_size=3, padding=1),
  12.             nn.ReLU(inplace=True),
  13.             nn.MaxPool2d(kernel_size=2, stride=2),
  14.             nn.Conv2d(128, 256, kernel_size=3, padding=1),
  15.             nn.ReLU(inplace=True),
  16.             nn.MaxPool2d(kernel_size=2, stride=2)
  17.         )
  18.         self.classifier = nn.Sequential(
  19.             nn.Linear(256 * 14 * 14, 1024),
  20.             nn.ReLU(inplace=True),
  21.             nn.Dropout(0.5),
  22.             nn.Linear(1024, 2)
  23.         )
  25.     def forward(self, x):
  26.         x = self.features(x)
  27.         x = x.view(x.size(0), -1)
  28.         x = self.classifier(x)
  29.         return x

迁移学习方案

推荐使用预训练的ResNet18模型进行迁移学习：

  1. def create_model(pretrained=True):
  2.     model = models.resnet18(pretrained=pretrained)
  3.     # 冻结所有卷积层参数
  4.     for param in model.parameters():
  5.         param.requires_grad = False
  6.     # 替换最后的全连接层
  7.     num_ftrs = model.fc.in_features
  8.     model.fc = nn.Sequential(
  9.         nn.Linear(num_ftrs, 512),
  10.         nn.ReLU(),
  11.         nn.Dropout(0.5),
  12.         nn.Linear(512, 2)
  13.     )
  14.     return model

训练流程优化

损失函数与优化器选择

  1. 交叉熵损失：nn.CrossEntropyLoss()
  2. 优化器：Adam（学习率0.001）或SGD with Momentum（学习率0.01，动量0.9）
  1. criterion = nn.CrossEntropyLoss()
  2. optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  3. scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

训练循环实现

  1. def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs=25):
  2.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  3.     model = model.to(device)
  5.     best_acc = 0.0
  6.     for epoch in range(num_epochs):
  7.         print(f'Epoch {epoch}/{num_epochs - 1}')
  8.         print('-' * 10)
  10.         # 每个epoch都有训练和验证阶段
  11.         for phase in ['train', 'val']:
  12.             if phase == 'train':
  13.                 model.train()  # 设置模型为训练模式
  14.             else:
  15.                 model.eval()   # 设置模型为评估模式
  17.             running_loss = 0.0
  18.             running_corrects = 0
  20.             # 迭代数据
  21.             for inputs, labels in dataloaders[phase]:
  22.                 inputs = inputs.to(device)
  23.                 labels = labels.to(device)
  25.                 # 梯度清零
  26.                 optimizer.zero_grad()
  28.                 # 前向传播
  29.                 with torch.set_grad_enabled(phase == 'train'):
  30.                     outputs = model(inputs)
  31.                     _, preds = torch.max(outputs, 1)
  32.                     loss = criterion(outputs, labels)
  34.                     # 反向传播+优化仅在训练阶段
  35.                     if phase == 'train':
  36.                         loss.backward()
  37.                         optimizer.step()
  39.                 # 统计
  40.                 running_loss += loss.item() * inputs.size(0)
  41.                 running_corrects += torch.sum(preds == labels.data)
  43.             if phase == 'train':
  44.                 scheduler.step()
  46.             epoch_loss = running_loss / len(dataloaders[phase].dataset)
  47.             epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
  49.             print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
  51.             # 深度复制模型
  52.             if phase == 'val' and epoch_acc > best_acc:
  53.                 best_acc = epoch_acc
  54.                 torch.save(model.state_dict(), 'best_model.pth')
  56.     print(f'Best val Acc: {best_acc:.4f}')
  57.     return model

预测与部署

测试集预测实现

  1. def predict_test(model, test_dir, transform):
  2.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  3.     model.eval()
  5.     test_images = []
  6.     for img_name in os.listdir(test_dir):
  7.         img_path = os.path.join(test_dir, img_name)
  8.         image = Image.open(img_path).convert('RGB')
  9.         if transform:
  10.             image = transform(image)
  11.         test_images.append((image, img_name))
  13.     predictions = []
  14.     with torch.no_grad():
  15.         for img, name in test_images:
  16.             img = img.unsqueeze(0).to(device)
  17.             outputs = model(img)
  18.             _, pred = torch.max(outputs, 1)
  19.             predictions.append((name, pred.item()))
  21.     return predictions

模型部署建议

  1. 导出为TorchScript格式：
    1. traced_script_module = torch.jit.trace(model, example_input)
    2. traced_script_module.save("model.pt")
  2. 使用ONNX格式跨平台部署：
    1. dummy_input = torch.randn(1, 3, 224, 224)
    2. torch.onnx.export(model, dummy_input, "model.onnx")

性能优化技巧

  1. 混合精度训练：使用torch.cuda.amp自动混合精度

    1. scaler = torch.cuda.amp.GradScaler()
    2. with torch.cuda.amp.autocast():
    3.  outputs = model(inputs)
    4.  loss = criterion(outputs, labels)
    5. scaler.scale(loss).backward()
    6. scaler.step(optimizer)
    7. scaler.update()

  2. 分布式训练：多GPU训练配置

    1. model = nn.DataParallel(model)
    2. model = model.to(device)

  3. 学习率预热：使用torch.optim.lr_scheduler.LambdaLR实现

常见问题解决方案

  1. 过拟合问题：

    • 增加Dropout层（建议0.3-0.5）
    • 使用L2正则化（weight_decay=0.001）
    • 增加数据增强强度

  2. 收敛缓慢问题：

    • 检查学习率是否合适（建议1e-3~1e-5）
    • 尝试不同的优化器（AdamW通常表现更好）
    • 检查数据预处理是否正确

  3. 内存不足问题：

    • 减小batch_size（推荐32-64）
    • 使用梯度累积（accumulate_grad_batches）
    • 清理中间变量（torch.cuda.empty_cache()）

完整代码示例

  1. # 主程序入口
  2. def main():
  3.     # 数据预处理
  4.     data_transforms = {
  5.         'train': transforms.Compose([
  6.             transforms.RandomResizedCrop(224),
  7.             transforms.RandomHorizontalFlip(),
  8.             transforms.ToTensor(),
  9.             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  10.         ]),
  11.         'val': transforms.Compose([
  12.             transforms.Resize(256),
  13.             transforms.CenterCrop(224),
  14.             transforms.ToTensor(),
  15.             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  16.         ]),
  17.     }
  19.     # 加载数据集
  20.     data_dir = 'data/dogs-vs-cats/train'
  21.     image_datasets = {x: CatDogDataset(os.path.join(data_dir, x),
  22.                                       data_transforms[x])
  23.                       for x in ['train', 'val']}
  24.     dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=32,
  25.                                                  shuffle=True, num_workers=4)
  26.                   for x in ['train', 'val']}
  28.     # 初始化模型
  29.     model = create_model(pretrained=True)
  31.     # 训练参数
  32.     criterion = nn.CrossEntropyLoss()
  33.     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  34.     scheduler = StepLR(optimizer, step_size=5, gamma=0.1)
  36.     # 训练模型
  37.     model = train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs=25)
  39.     # 保存模型
  40.     torch.save(model.state_dict(), 'cat_dog_classifier.pth')
  42. if __name__ == '__main__':
  43.     main()

通过以上完整实现，开发者可以系统掌握使用PyTorch完成图像分类任务的全流程。实际项目中建议：1）优先使用迁移学习方案；2）重视数据增强策略；3）合理设置学习率调度；4）监控验证集性能防止过拟合。该方案在Kaggle测试集上可达到98%以上的准确率。

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数