3. 数据准备：图像分类数据集下载与组织

做 MobileViT 训练，第一步就是搞定数据。说实话，很多人觉得这步没啥技术含量，随便找个数据集就开干。但我见过太多项目，最后模型训不出来，回头一查，是数据路径写错了，或者类别标签对不上。嗯，这种坑我踩过不止一次。

今天咱们就把数据准备这件事，一次性说清楚。从数据集下载、目录结构组织，到 DataLoader 编写，一条龙搞定。

3.1 数据集选择：ImageNet 子集 vs 自定义数据集

我个人习惯把数据集分成两类：一类是拿来练手的公开数据集，另一类是实际项目里的自定义数据。

ImageNet 子集：完整 ImageNet 有 1000 类、1400 多万张图，太大了。咱们做 MobileViT 实验，通常用它的子集。比如只取 10 类或 100 类，每类 500-1000 张图。这样训练快，调试也方便。

自定义数据集：实际项目中，你拿到的数据往往长这样——一堆图片散落在各个文件夹里，命名乱七八糟。这时候就需要我们手动整理。

核心原则：不管什么数据集，最终都要组织成 ImageNet 风格的目录结构。这是 PyTorch 的 ImageFolder 默认支持的格式，省心。

3.2 目录结构组织：ImageNet 风格

先看标准结构长什么样：

dataset/
├── train/
│   ├── class_0/
│   │   ├── img_001.jpg
│   │   ├── img_002.jpg
│   │   └── ...
│   ├── class_1/
│   │   ├── img_001.jpg
│   │   └── ...
│   └── ...
└── val/
    ├── class_0/
    │   ├── img_001.jpg
    │   └── ...
    ├── class_1/
    │   └── ...
    └── ...

每个类别一个文件夹，文件夹名就是类别名。训练集和验证集分开。就这么简单。

我曾经接手过一个项目，对方把训练集和验证集混在一起，还用了中文文件夹名。结果 PyTorch 的 ImageFolder 直接报错，因为默认编码不支持中文路径。后来我写了个脚本批量重命名，折腾了一下午。

避坑指南：文件夹名最好用英文，不要带空格。比如用 "golden_retriever" 而不是 "Golden Retriever"。Linux 下空格处理起来很麻烦。

3.3 自定义数据集整理脚本

如果你手头的数据是散乱的，可以用下面这个脚本快速整理：

import os
import shutil
from sklearn.model_selection import train_test_split

def organize_dataset(src_dir, dst_dir, val_ratio=0.2):
    """
    将散乱的数据整理成 ImageNet 风格
    src_dir: 原始数据目录，每个子文件夹是一个类别
    dst_dir: 目标目录
    """
    # 创建目标目录
    train_dir = os.path.join(dst_dir, 'train')
    val_dir = os.path.join(dst_dir, 'val')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    
    # 遍历每个类别
    for class_name in os.listdir(src_dir):
        class_path = os.path.join(src_dir, class_name)
        if not os.path.isdir(class_path):
            continue
            
        # 获取该类别的所有图片
        images = [f for f in os.listdir(class_path) 
                  if f.endswith(('.jpg', '.jpeg', '.png'))]
        
        # 划分训练集和验证集
        train_imgs, val_imgs = train_test_split(
            images, test_size=val_ratio, random_state=42
        )
        
        # 创建类别子目录
        os.makedirs(os.path.join(train_dir, class_name), exist_ok=True)
        os.makedirs(os.path.join(val_dir, class_name), exist_ok=True)
        
        # 复制文件
        for img in train_imgs:
            shutil.copy(
                os.path.join(class_path, img),
                os.path.join(train_dir, class_name, img)
            )
        for img in val_imgs:
            shutil.copy(
                os.path.join(class_path, img),
                os.path.join(val_dir, class_name, img)
            )
        
        print(f"类别 {class_name}: 训练集 {len(train_imgs)} 张, 验证集 {len(val_imgs)} 张")

# 使用示例
organize_dataset('./raw_data', './dataset', val_ratio=0.2)

这个脚本我用了好几年，基本没出过问题。唯一要注意的是，如果你的数据量特别大（比如几十万张），用 shutil.copy 会慢，可以改成 os.symlink 创建软链接，省空间也省时间。

3.4 DataLoader 编写：从基础到优化

数据准备好了，接下来就是写 DataLoader。这里我直接给出一套完整的代码，包含了 MobileViT 训练时常用的数据增强策略。

3.4.1 基础版 DataLoader

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

def get_dataloaders(data_dir, batch_size=64, num_workers=4):
    """
    获取训练和验证的 DataLoader
    """
    # 训练集：带数据增强
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        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])
    ])
    
    # 加载数据集
    train_dataset = datasets.ImageFolder(
        os.path.join(data_dir, 'train'),
        transform=train_transform
    )
    val_dataset = datasets.ImageFolder(
        os.path.join(data_dir, 'val'),
        transform=val_transform
    )
    
    # 创建 DataLoader
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, val_loader, train_dataset.classes

注意：num_workers 不是越大越好。我试过在 8 核机器上设成 8，结果 CPU 被打满，反而比设成 4 还慢。一般建议设成 CPU 核心数的一半。

3.4.2 进阶版：MobileViT 专用增强

MobileViT 对输入尺寸比较敏感。我个人建议在训练时用 RandomResizedCrop 配合 RandAugment，效果比单纯翻转好很多。

from torchvision.transforms import RandAugment

# MobileViT 推荐的数据增强
mobilevit_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.08, 1.0)),
    transforms.RandomHorizontalFlip(),
    RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
])

为什么用 scale=(0.08, 1.0)？因为 MobileViT 的 patch embedding 层对图像内容占比有要求。如果裁剪得太小，patch 里可能全是背景，学不到有用特征。这个参数是我在 ImageNet 上反复试出来的。

3.5 核心逻辑流程图

下面这张图展示了数据准备的完整流程，从原始数据到 DataLoader 的每一步：

数据准备核心流程 原始数据 散乱图片 / 公开数据集 目录结构组织 ImageNet 风格：train/class/ val/class 数据增强 RandomResizedCrop + RandAugment DataLoader batch_size / num_workers / pin_memory 训练就绪 关键点 类别名用英文 技巧 scale 设 0.08~1.0 避坑 num_workers 别太大

3.6 常见问题与解决方案

问题 原因 解决方案
DataLoader 卡住不动 num_workers 设置过大，或内存不足 减少 num_workers，或设置 prefetch_factor=2
训练时内存暴涨 pin_memory=True 但没配合非阻塞传输 在训练循环中加 with torch.no_grad()
类别标签不对 文件夹排序和预期不一致 打印 dataset.class_to_idx 检查映射关系
图片加载失败 图片损坏或格式不支持 用 PIL.Image.open() 逐个检查，删除坏图

我的经验：每次换数据集，我都会先跑一个 mini-batch 看看图片和标签对不对应。花 5 分钟检查，能省下后面 5 小时的 debug 时间。

好了，数据准备这块就讲到这里。你按照上面的步骤来，基本不会出大问题。记住：数据质量决定了模型性能的上限，花时间把数据整理好，值得。

可可精品资料库

← 上一章 📖 返回目录 下一章 →