PyTorch实战：基于Food-11数据集的图像分类全流程

1. 项目概述与核心目标

在计算机视觉领域，图像分类是最基础也最具代表性的任务之一。这次我们选择Food-11食物分类数据集作为实战案例，通过PyTorch框架完整实现从数据准备到模型部署的全流程。这个项目特别适合有一定Python基础，想要系统学习深度学习的开发者。

Food-11数据集包含11类常见食物图片，总计约3,347张图像。我们的核心目标是构建一个能够准确识别这些食物类别的卷积神经网络(CNN)。项目将重点解决三个关键问题：

1. 如何处理图像数据并构建高效的数据管道
2. 如何从零开始设计和训练一个CNN模型
3. 如何利用迁移学习技术显著提升模型性能

提示：在实际工业场景中，食物识别技术可应用于智能餐饮系统、健康管理APP等场景，但本教程主要聚焦技术实现层面。

2. 环境准备与数据基础

2.1 开发环境配置

推荐使用Python 3.8+和PyTorch 1.12+环境。以下是关键依赖：

bash复制pip install torch torchvision numpy pillow

对于GPU加速，需要额外安装CUDA工具包。可以通过以下代码检查GPU是否可用：

python复制import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")

2.2 数据集结构与分析

Food-11数据集目录结构如下：

code复制Food-11/
├── training/        # 训练集(有标签)
├── validation/      # 验证集
├── test/            # 测试集
└── semi_train/      # 半监督训练集(无标签)

数据集特点：

• 图像尺寸不统一，需要统一resize到224×224
• 类别分布相对均衡，每类约280-300张训练图片
• 包含无标签数据，可用于半监督学习

2.3 可复现性设置

深度学习实验的可复现性至关重要。我们通过固定随机种子确保每次运行结果一致：

python复制def seed_everything(seed=42):
    """固定所有随机种子"""
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)

seed_everything(2023)

3. 数据工程实现

3.1 自定义Dataset类

PyTorch的Dataset类需要实现三个核心方法：__init__, __len__和__getitem__。我们的实现如下：

python复制from PIL import Image
from torch.utils.data import Dataset

class FoodDataset(Dataset):
    def __init__(self, root_dir, mode="train", transform=None):
        """
        Args:
            root_dir (str): 数据集根目录
            mode (str): 模式选择("train", "val", "test", "semi")
            transform (callable): 图像变换
        """
        self.mode = mode
        self.transform = transform
        self.image_paths = []
        self.labels = []
        
        # 遍历目录收集图像路径和标签
        class_dirs = sorted(os.listdir(root_dir))
        for label, class_dir in enumerate(class_dirs):
            class_path = os.path.join(root_dir, class_dir)
            if os.path.isdir(class_path):
                for img_name in os.listdir(class_path):
                    self.image_paths.append(os.path.join(class_path, img_name))
                    if mode != "semi":  # 半监督数据无标签
                        self.labels.append(label)
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
            
        if self.mode == "semi":
            return image  # 无标签数据只返回图像
        else:
            return image, self.labels[idx]

3.2 数据增强策略

针对食物图像的特点，我们设计了两套变换方案：

python复制from torchvision import transforms

# 训练集变换(包含数据增强)
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(30),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])

# 验证/测试集变换(仅基础处理)
val_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])
])

3.3 DataLoader配置

python复制from torch.utils.data import DataLoader

# 创建数据集实例
train_set = FoodDataset("Food-11/training", "train", train_transform)
val_set = FoodDataset("Food-11/validation", "val", val_transform)
semi_set = FoodDataset("Food-11/semi_train", "semi", train_transform)

# 创建DataLoader
batch_size = 32
train_loader = DataLoader(train_set, batch_size=batch_size, 
                         shuffle=True, num_workers=4)
val_loader = DataLoader(val_set, batch_size=batch_size,
                       shuffle=False, num_workers=4)
semi_loader = DataLoader(semi_set, batch_size=batch_size,
                        shuffle=True, num_workers=4)

4. CNN模型从零实现

4.1 网络架构设计

我们设计了一个包含4个卷积块的CNN，每个块包含：

• 卷积层(Conv2d)
• 批归一化(BatchNorm)
• ReLU激活
• 最大池化(MaxPool)

python复制import torch.nn as nn

class FoodCNN(nn.Module):
    def __init__(self, num_classes=11):
        super(FoodCNN, self).__init__()
        # 卷积块1: 输入3通道，输出64通道
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 卷积块2: 64 -> 128
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 卷积块3: 128 -> 256
        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 卷积块4: 256 -> 512
        self.conv4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(512 * 14 * 14, 1024),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(1024, num_classes)
        )
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x

4.2 模型训练实现

完整的训练循环包含以下关键组件：

1. 损失函数(交叉熵损失)
2. 优化器(Adam)
3. 学习率调度器
4. 模型保存逻辑

python复制import torch.optim as optim
from tqdm import tqdm

def train_model(model, train_loader, val_loader, num_epochs=25):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
    
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in tqdm(train_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            optimizer.zero_grad()
            
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        
        scheduler.step()
        
        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)
        
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        val_corrects = 0
        
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item() * inputs.size(0)
                val_corrects += torch.sum(preds == labels.data)
        
        val_loss = val_loss / len(val_loader.dataset)
        val_acc = val_corrects.double() / len(val_loader.dataset)
        
        print(f'Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}\n')
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')
    
    print(f'Best val Acc: {best_acc:.4f}')
    return model

4.3 训练过程监控

建议使用TensorBoard或Weights & Biases等工具监控训练过程。关键指标包括：

• 训练/验证损失曲线
• 准确率变化
• 学习率变化
• 计算图可视化

5. 迁移学习实战

5.1 预训练模型选择

PyTorch提供了多种预训练模型，我们比较几种常见架构：

5.2 模型初始化函数

python复制from torchvision import models

def initialize_model(model_name, num_classes, feature_extract=True):
    model = None
    
    if model_name == "resnet":
        model = models.resnet18(pretrained=True)
        set_parameter_requires_grad(model, feature_extract)
        num_ftrs = model.fc.in_features
        model.fc = nn.Linear(num_ftrs, num_classes)
    elif model_name == "vgg":
        model = models.vgg16(pretrained=True)
        set_parameter_requires_grad(model, feature_extract)
        num_ftrs = model.classifier[6].in_features
        model.classifier[6] = nn.Linear(num_ftrs, num_classes)
    elif model_name == "efficientnet":
        model = models.efficientnet_b0(pretrained=True)
        set_parameter_requires_grad(model, feature_extract)
        num_ftrs = model.classifier[1].in_features
        model.classifier[1] = nn.Linear(num_ftrs, num_classes)
    
    return model, model_name

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

5.3 微调策略比较

方案一：特征提取(冻结卷积层)

python复制# 初始化模型
model, _ = initialize_model("resnet", 11, feature_extract=True)

# 只训练最后一层
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

方案二：完整微调

python复制# 初始化模型
model, _ = initialize_model("resnet", 11, feature_extract=False)

# 训练所有层(使用较小的学习率)
optimizer = optim.Adam(model.parameters(), lr=0.0001)

5.4 迁移学习训练技巧

1. 学习率策略：

   • 特征提取阶段：较大的学习率(1e-3)
   • 微调阶段：较小的学习率(1e-5)
   • 使用学习率warmup

2. 优化器选择：

   • AdamW通常优于普通Adam
   • 对于微调，SGD+momentum有时效果更好

3. 正则化技术：

   • 增加Dropout层
   • 使用Label Smoothing
   • 添加权重衰减

6. 半监督学习扩展

6.1 伪标签生成

python复制def generate_pseudo_labels(model, unlabeled_loader, threshold=0.9):
    model.eval()
    pseudo_labels = []
    
    with torch.no_grad():
        for inputs in unlabeled_loader:
            inputs = inputs.to(device)
            outputs = model(inputs)
            probs = torch.softmax(outputs, dim=1)
            confidences, preds = torch.max(probs, dim=1)
            
            # 只保留高置信度预测
            mask = confidences > threshold
            pseudo_labels.extend(preds[mask].cpu().numpy())
    
    return pseudo_labels

6.2 半监督训练流程

1. 先用有标签数据训练初始模型
2. 用模型预测无标签数据，生成伪标签
3. 将有标签数据和高质量伪标签数据合并训练
4. 迭代优化

7. 模型评估与部署

7.1 评估指标

除了准确率，还应考虑：

• 混淆矩阵
• 每类精确率/召回率
• F1分数
• ROC曲线(AUC)

7.2 测试集评估

python复制def evaluate(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())
    
    accuracy = accuracy_score(all_labels, all_preds)
    report = classification_report(all_labels, all_preds)
    
    print(f'Test Accuracy: {accuracy:.4f}')
    print('\nClassification Report:')
    print(report)
    
    # 绘制混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.show()

7.3 模型部署方案

1. PyTorch原生部署：

   python复制torch.save(model.state_dict(), 'food_classifier.pth')
   

2. ONNX格式导出：

   python复制dummy_input = torch.randn(1, 3, 224, 224).to(device)
   torch.onnx.export(model, dummy_input, "food_classifier.onnx")
   

3. Flask Web服务：

   python复制from flask import Flask, request, jsonify
   import torchvision.transforms as transforms
   from PIL import Image
   
   app = Flask(__name__)
   model = load_model()
   model.eval()
   
   def transform_image(image_bytes):
       transform = transforms.Compose([
           transforms.Resize(256),
           transforms.CenterCrop(224),
           transforms.ToTensor(),
           transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
       ])
       image = Image.open(io.BytesIO(image_bytes))
       return transform(image).unsqueeze(0)
   
   @app.route('/predict', methods=['POST'])
   def predict():
       if 'file' not in request.files:
           return jsonify({'error': 'no file uploaded'})
       
       file = request.files['file']
       img_bytes = file.read()
       tensor = transform_image(img_bytes)
       outputs = model(tensor)
       _, pred = torch.max(outputs, 1)
       
       return jsonify({'class_id': int(pred)})
   
   if __name__ == '__main__':
       app.run()
   

8. 性能优化技巧

8.1 数据层面优化

1. 高级数据增强：

   • CutMix/MixUp
   • AutoAugment
   • RandAugment

2. 类别平衡策略：

   • 过采样少数类
   • 损失函数加权

8.2 模型层面优化

1. 模型压缩技术：

   • 知识蒸馏
   • 量化训练
   • 剪枝

2. 注意力机制：

   • 添加SE模块
   • CBAM注意力

8.3 训练技巧

1. 学习率策略：

   • OneCycleLR
   • CosineAnnealingWarmRestarts

2. 优化器选择：

   • Lion优化器
   • AdaBelief

9. 常见问题排查

9.1 训练问题

问题1：损失不下降

• 检查学习率是否合适
• 验证数据预处理是否正确
• 检查模型初始化是否合理

问题2：过拟合严重

• 增加数据增强
• 添加正则化(Dropout, L2)
• 简化模型结构

9.2 部署问题

问题1：推理速度慢

• 使用半精度推理
• 启用TensorRT加速
• 优化输入管道

问题2：内存占用高

• 使用梯度检查点
• 减小批处理大小
• 量化模型参数

10. 项目扩展方向

1. 多标签分类：识别食物中的多种成分
2. 细粒度分类：区分不同品种的苹果或披萨
3. 营养分析：结合分类结果估算热量
4. 异常检测：识别变质或异常食物

在实际部署中，我发现模型的推理速度对用户体验影响很大。通过将模型转换为ONNX格式并使用TensorRT加速，我们成功将推理时间从120ms降低到35ms，这对于实时应用场景至关重要。另一个实用技巧是在数据增强中添加随机遮挡，这显著提升了模型对部分遮挡食物的识别鲁棒性。