从零实现食物分类：小数据集深度学习实战

1. 项目概述

作为一名长期奋战在计算机视觉一线的算法工程师，我深知图像分类任务在深度学习领域的基础性和重要性。今天，我将以food-11食物分类项目为例，带大家完整走一遍从零搭建分类模型的全流程。这个项目虽然看似简单，但包含了数据预处理、模型设计、训练优化等核心环节，是理解深度学习图像处理的绝佳案例。

在工业界实践中，我们经常遇到类似food-11这样的小规模数据集场景。这类项目的特点是：数据量有限（通常每个类别只有几十到几百张图片）、类别间差异可能不明显（比如不同种类的面食）、但业务需求又要求较高的准确率。经过多次实战，我总结出了一套针对小规模数据集的解决方案，今天就将这些经验毫无保留地分享给大家。

2. 数据准备与预处理

2.1 数据集结构解析

food-11数据集采用典型的分类数据集结构，包含三个子集：

code复制food-11/
├── training/
│   ├── labeled/   # 带标注的训练数据
│   └── unlabeled/ # 无标注数据（可用于半监督学习）
├── validation/    # 验证集
└── testing/       # 测试集

每个类别文件夹以两位数编号命名（00-10），这种结构在实际项目中非常常见。我建议在开始编码前，先用tree命令或文件管理器直观查看数据集结构，这能帮助我们更好地设计数据读取逻辑。

2.2 数据读取实现

数据读取是深度学习项目中最容易被忽视但极其重要的环节。一个健壮的数据读取管道应该具备以下特性：

• 内存高效（特别是处理大型数据集时）
• 支持多种图像格式
• 具备基本的数据校验功能

以下是改进后的数据读取实现：

python复制def read_file(path, img_size=224):
    """
    读取分类数据集
    :param path: 数据集根路径
    :param img_size: 统一调整的图像尺寸
    :return: 图像数据(numpy数组)和标签
    """
    X, Y = [], []
    class_dirs = sorted([d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))])
    
    for class_idx, class_dir in enumerate(tqdm(class_dirs)):
        class_path = os.path.join(path, class_dir)
        img_files = [f for f in os.listdir(class_path) 
                    if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        for img_file in img_files:
            try:
                img_path = os.path.join(class_path, img_file)
                img = Image.open(img_path).convert('RGB')  # 确保转为RGB格式
                img = img.resize((img_size, img_size))
                X.append(np.array(img))
                Y.append(class_idx)
            except Exception as e:
                print(f"Error loading {img_path}: {str(e)}")
                continue
                
    return np.array(X), np.array(Y)

这个实现有以下改进：

1. 自动过滤非图像文件
2. 强制转换为RGB格式（避免某些灰度图像导致维度问题）
3. 添加了异常处理
4. 使用列表存储再转换，比预分配数组更灵活

2.3 数据增强策略

对于小数据集，数据增强是提升模型泛化能力的关键。在food分类任务中，我们需要考虑食物图像的特殊性：

1. 几何变换：

   • 随机旋转（-30°到30°）：模拟不同拍摄角度
   • 随机水平翻转：适用于对称性较强的食物
   • 随机裁剪：模拟局部特写

2. 颜色变换：

   • 亮度/对比度调整：模拟不同光照条件
   • 色相/饱和度微调：考虑食物颜色的多样性

python复制from torchvision import transforms

train_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.RandomRotation(30),
    transforms.RandomHorizontalFlip(p=0.5),
    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.ToPILImage(),
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

重要细节：

• 验证集使用CenterCrop而不是RandomCrop，保证评估一致性
• 标准化参数使用ImageNet的均值和标准差，这对使用预训练模型很重要
• RandomResizedCrop的scale参数控制裁剪比例，0.8-1.0是个经验值

3. 模型设计与实现

3.1 自定义CNN模型

对于初学者来说，从零搭建CNN是理解卷积神经网络工作原理的最佳方式。我们的自定义模型遵循经典设计范式：

python复制class FoodCNN(nn.Module):
    def __init__(self, num_classes=11):
        super(FoodCNN, self).__init__()
        # 特征提取器
        self.features = nn.Sequential(
            nn.Conv2d(3, 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),
            
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # 分类器
        self.classifier = nn.Sequential(
            nn.Linear(512 * 14 * 14, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(1024, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # flatten
        x = self.classifier(x)
        return x

设计要点：

1. 采用渐进式下采样策略，逐步增加通道数、减小空间尺寸
2. 每个卷积层后接BatchNorm和ReLU，这是现代CNN的标准配置
3. 分类器部分加入Dropout防止过拟合
4. 特征图尺寸计算：(224/2^4)=14，所以最后是512x14x14

3.2 训练流程实现

一个完整的训练流程需要包含以下组件：

• 损失函数（交叉熵损失）
• 优化器（AdamW）
• 学习率调度器
• 训练-验证循环

python复制def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, scheduler=None):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs-1}')
        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 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)
            
            if phase == 'train' and scheduler:
                scheduler.step()
            
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
            
            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())
    
    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed//60:.0f}m {time_elapsed%60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')
    
    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

关键细节：

1. 使用torch.set_grad_enabled控制是否计算梯度，提高验证阶段效率
2. 记录最佳模型权重，避免过拟合导致模型退化
3. 添加学习率调度器支持（如CosineAnnealingLR）
4. 使用double类型计算准确率，避免整数除法问题

4. 迁移学习实战

4.1 为什么需要迁移学习

在小数据集场景下（如food-11），从零训练CNN模型通常会遇到两个问题：

1. 模型容易过拟合，因为参数太多而数据太少
2. 训练不充分，难以收敛到好的解

迁移学习通过使用在大规模数据集（如ImageNet）上预训练的模型，可以显著提升小数据集的分类性能。根据经验，使用迁移学习通常能带来20-50%的准确率提升。

4.2 ResNet18微调实战

python复制from torchvision import models

def setup_resnet(num_classes=11, pretrained=True):
    # 加载预训练模型
    model = models.resnet18(pretrained=pretrained)
    
    # 冻结所有卷积层参数
    for param in model.parameters():
        param.requires_grad = False
    
    # 替换最后的全连接层
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)
    
    # 只对新添加的分类层设置requires_grad=True
    for param in model.fc.parameters():
        param.requires_grad = True
        
    return model

# 初始化模型
model = setup_resnet(num_classes=11).to(device)

# 只对需要梯度的参数进行优化
optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3,
    weight_decay=1e-4
)

微调策略：

1. 初始阶段冻结所有卷积层，只训练最后的分类层（特征提取器模式）
2. 训练几轮后，可以解冻部分高层卷积层进行微调
3. 使用较小的学习率（通常比从零训练小10倍）

4.3 不同模型的性能对比

我在food-11验证集上对比了几种常见架构的表现：

从结果可以看出，即使是小型的EfficientNet-B0，也能在参数量较少的情况下取得不错的性能。在实际项目中，我们需要在准确率和推理速度之间做出权衡。

5. 工程化改进

5.1 模型封装工厂

为了提高代码复用性，我们可以实现一个模型工厂：

python复制class ModelFactory:
    @staticmethod
    def create_model(model_name, num_classes, pretrained=True):
        model_map = {
            'resnet18': models.resnet18,
            'resnet50': models.resnet50,
            'efficientnet': models.efficientnet_b0,
            'custom': lambda pretrained: FoodCNN(num_classes)
        }
        
        if model_name not in model_map:
            raise ValueError(f"Unsupported model: {model_name}")
            
        model = model_map[model_name](pretrained=pretrained)
        
        # 替换分类头
        if model_name.startswith('resnet'):
            num_ftrs = model.fc.in_features
            model.fc = nn.Linear(num_ftrs, num_classes)
        elif model_name.startswith('efficientnet'):
            num_ftrs = model.classifier[1].in_features
            model.classifier[1] = nn.Linear(num_ftrs, num_classes)
            
        return model

5.2 训练可视化

使用TensorBoard记录训练过程：

python复制from torch.utils.tensorboard import SummaryWriter

def train_with_logging(...):
    writer = SummaryWriter()
    
    for epoch in range(num_epochs):
        # ... 训练逻辑 ...
        
        # 记录标量
        writer.add_scalar('Loss/train', epoch_loss, epoch)
        writer.add_scalar('Accuracy/train', epoch_acc, epoch)
        
        # 记录直方图
        for name, param in model.named_parameters():
            writer.add_histogram(name, param, epoch)
    
    writer.close()

5.3 超参数优化

使用Optuna进行自动化超参数搜索：

python复制import optuna

def objective(trial):
    lr = trial.suggest_float('lr', 1e-5, 1e-3, log=True)
    weight_decay = trial.suggest_float('weight_decay', 1e-6, 1e-3)
    batch_size = trial.suggest_categorical('batch_size', [16, 32, 64])
    
    # 初始化模型和数据加载器
    model = ModelFactory.create_model('resnet18', 11)
    train_loader, val_loader = create_dataloaders(batch_size)
    
    # 训练过程
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    train_model(model, {'train': train_loader, 'val': val_loader}, optimizer=optimizer)
    
    # 返回验证集准确率作为优化目标
    return evaluate(model, val_loader)

study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=50)
print(f'Best trial: {study.best_trial.params}')

6. 部署优化建议

当模型训练完成后，我们需要考虑如何部署到生产环境。以下是几个实用建议：

1. 模型量化：使用torch.quantization减小模型大小，提升推理速度

   python复制quantized_model = torch.quantization.quantize_dynamic(
       model, {nn.Linear}, dtype=torch.qint8
   )
   

2. ONNX导出：转换为通用格式便于跨平台部署

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

3. TorchScript序列化：提升Python推理性能

   python复制scripted_model = torch.jit.script(model)
   scripted_model.save("food_classifier.pt")
   

4. Web服务封装：使用FastAPI创建REST API

   python复制from fastapi import FastAPI, File, UploadFile
   app = FastAPI()
   
   @app.post("/predict")
   async def predict(image: UploadFile = File(...)):
       img = Image.open(image.file).convert('RGB')
       tensor = val_transform(img).unsqueeze(0)
       with torch.no_grad():
           outputs = model(tensor)
       _, pred = torch.max(outputs, 1)
       return {"class": class_names[pred.item()]}
   

7. 项目扩展方向

这个基础项目可以进一步扩展为更实用的应用：

1. 多标签分类：有些食物可能属于多个类别（如"辣"+"面食"）

   • 使用sigmoid输出代替softmax
   • 采用Binary Cross Entropy损失

2. 细粒度分类：区分更细的食物类别（如不同种类的苹果）

   • 使用更强大的backbone（如ResNet152）
   • 添加注意力机制

3. 营养分析：结合分类结果估算食物热量

   • 构建食物-营养数据库
   • 开发回归模型预测热量

4. 移动端部署：开发手机应用

   • 使用TensorFlow Lite或Core ML
   • 优化模型大小和推理速度

在实际开发中，我发现以下几个技巧特别有用：

1. 使用混合精度训练可以显著减少显存占用
2. 对图像分类任务，适当的数据增强比模型架构更重要
3. 当准确率停滞时，尝试调整学习率比增加epoch更有效
4. 模型集成（如3-5个不同架构的模型）通常能带来2-5%的提升