PyTorch实现运动鞋图像分类：从CNN到迁移学习

1. 项目概述

运动鞋识别是计算机视觉领域的一个典型应用场景。通过深度学习技术，我们可以训练一个能够自动识别不同品牌或款式运动鞋的模型。这个项目使用PyTorch框架构建了一个卷积神经网络(CNN)，实现了对运动鞋图像的分类任务。

在实际应用中，运动鞋识别技术可以用于电商平台的商品搜索、库存管理、防伪鉴定等多个场景。例如，用户可以通过拍照上传的方式快速找到同款运动鞋，或者验证所购商品的真伪。

提示：本项目采用了经典的深度学习流程，包括数据准备、模型构建、训练优化和结果评估四个主要环节。即使你是深度学习新手，跟随这个流程也能快速上手图像分类任务。

2. 前期准备工作

2.1 环境配置与数据准备

首先需要搭建PyTorch开发环境。推荐使用Anaconda创建Python虚拟环境：

bash复制conda create -n shoe_recognition python=3.8
conda activate shoe_recognition
pip install torch torchvision torchaudio

数据集组织是项目成功的关键。建议按照以下目录结构存放运动鞋图片：

code复制5-data/
    ├── train/
    │   ├── brand1/
    │   ├── brand2/
    │   └── ...
    └── test/
        ├── brand1/
        ├── brand2/
        └── ...

每个子目录代表一个类别，存放对应品牌的运动鞋图片。这种结构可以直接被PyTorch的ImageFolder加载器使用。

2.2 数据预处理与增强

数据预处理对模型性能影响显著。我们定义了两组变换：

python复制train_transforms = transforms.Compose([
    transforms.Resize([224, 224]),  # 统一图像尺寸
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 颜色扰动
    transforms.ToTensor(),  # 转为张量
    transforms.Normalize(  # 标准化
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225])
])

注意：训练集和测试集必须使用相同的标准化参数。这里使用的是ImageNet的均值和标准差，因为它们已经在大量图像上统计得到，适合大多数视觉任务。

2.3 数据加载与批处理

使用DataLoader可以高效地加载和批处理数据：

python复制batch_size = 32
train_dl = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,  # 训练集需要打乱
    num_workers=4  # 多线程加速加载
)

test_dl = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,  # 测试集不需要打乱
    num_workers=4
)

3. 构建CNN网络

3.1 基础CNN架构设计

我们首先构建一个简单的CNN网络：

python复制class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        
        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(128 * 28 * 28, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

这个网络包含三个卷积块，每个块由卷积层、批归一化、ReLU激活和最大池化组成。最后的分类器使用两个全连接层，中间加入Dropout防止过拟合。

3.2 网络结构优化技巧

在实际应用中，我们发现几个优化点：

1. 使用更深的网络结构：增加网络深度可以提取更高级的特征，但要注意梯度消失问题
2. 引入残差连接：帮助梯度回传，训练更深的网络
3. 全局平均池化替代全连接层：减少参数数量，防止过拟合

优化后的网络结构：

python复制class AdvancedCNN(nn.Module):
    def __init__(self, num_classes):
        super(AdvancedCNN, self).__init__()
        
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256, num_classes)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

4. 模型训练与优化

4.1 训练流程实现

训练循环的核心代码如下：

python复制def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, labels in dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    
    epoch_loss = running_loss / len(dataloader)
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

4.2 学习率调整策略

合适的学习率对训练至关重要。我们实现了多种学习率调整策略：

1. 固定学习率衰减：

python复制scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

2. 基于验证集性能的自适应调整：

python复制scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='max', factor=0.5, patience=3
)

3. 余弦退火学习率：

python复制scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=20, eta_min=1e-6
)

4.3 模型评估与验证

验证集评估代码：

python复制def validate(model, dataloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    val_loss = running_loss / len(dataloader)
    val_acc = correct / total
    return val_loss, val_acc

5. 迁移学习实践

5.1 使用预训练模型

对于小数据集，使用预训练模型可以显著提升性能：

python复制class ShoeClassifier(nn.Module):
    def __init__(self, num_classes):
        super(ShoeClassifier, self).__init__()
        self.backbone = models.resnet18(pretrained=True)
        
        # 冻结前几层参数
        for param in self.backbone.parameters():
            param.requires_grad = False
        
        # 替换最后的全连接层
        num_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(num_features, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)

5.2 分层学习率设置

不同层可以使用不同的学习率：

python复制optimizer = torch.optim.Adam([
    {'params': model.backbone.parameters(), 'lr': 1e-4},
    {'params': model.backbone.fc.parameters(), 'lr': 1e-3}
])

6. 结果分析与可视化

6.1 训练过程可视化

使用Matplotlib绘制训练曲线：

python复制plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_accs, label='Train Accuracy')
plt.plot(val_accs, label='Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

6.2 混淆矩阵分析

评估模型在各个类别上的表现：

python复制from sklearn.metrics import confusion_matrix
import seaborn as sns

def plot_confusion_matrix(model, dataloader, class_names, device):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            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=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.show()

7. 模型部署与应用

7.1 模型保存与加载

训练完成后保存最佳模型：

python复制torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'class_to_idx': train_dataset.class_to_idx
}, 'best_model.pth')

加载保存的模型：

python复制checkpoint = torch.load('best_model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
class_to_idx = checkpoint['class_to_idx']
idx_to_class = {v: k for k, v in class_to_idx.items()}

7.2 单张图片预测

实现预测函数：

python复制def predict_image(image_path, model, transform, idx_to_class, device):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    
    model.eval()
    with torch.no_grad():
        output = model(image)
        _, pred = torch.max(output, 1)
        class_id = pred.item()
        probabilities = torch.softmax(output, dim=1)[0] * 100
    
    return idx_to_class[class_id], probabilities[class_id].item()

8. 性能优化技巧

8.1 混合精度训练

使用AMP加速训练：

python复制scaler = torch.cuda.amp.GradScaler()

for inputs, labels in train_loader:
    inputs = inputs.to(device)
    labels = labels.to(device)
    
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

8.2 数据加载优化

使用更高效的数据加载方式：

python复制train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True
)

9. 常见问题与解决方案

9.1 过拟合问题

解决方案：

1. 增加数据增强
2. 使用更严格的Dropout
3. 添加权重衰减
4. 使用早停策略

9.2 训练不收敛

可能原因：

1. 学习率设置不当
2. 数据预处理有问题
3. 模型初始化不当

解决方案：

1. 尝试不同的学习率
2. 检查数据预处理流程
3. 使用预训练模型

9.3 类别不平衡

处理方法：

1. 使用加权损失函数
2. 过采样少数类
3. 欠采样多数类

python复制class_counts = get_class_counts(train_dataset)
class_weights = 1. / torch.tensor(class_counts, dtype=torch.float)
class_weights = class_weights / class_weights.sum()
criterion = nn.CrossEntropyLoss(weight=class_weights.to(device))

10. 项目扩展方向

1. 多标签分类：识别运动鞋的多个属性（品牌、款式、颜色等）
2. 细粒度分类：区分同一品牌的不同型号
3. 目标检测：定位图像中的运动鞋位置
4. 相似度匹配：找到与查询图片相似的运动鞋
5. 3D形状分析：从多视角图像重建运动鞋3D模型

在实际项目中，我发现使用ResNet18预训练模型配合适当的数据增强，可以在运动鞋识别任务上达到90%以上的准确率。关键是要根据具体需求平衡模型复杂度和计算资源，同时注意数据质量对最终性能的决定性影响。