CIFAR100数据集实战指南：从数据加载到模型训练的最佳实践

1次阅读

共计 3588 个字符，预计需要花费 9 分钟才能阅读完成。

CIFAR100 作为经典的图像分类基准数据集，在实际使用中常遇到数据预处理效率低、类别不均衡等问题。本文从 PyTorch 数据加载优化入手，详解如何高效处理 CIFAR100 数据集，包括数据增强策略、内存优化技巧，并给出完整的 ResNet 训练示例。读者将掌握工业级图像分类任务的数据处理全流程，获得 2 - 3 倍训练速度提升。

CIFAR100 数据集包含 60,000 张 32×32 像素的彩色图像，分为 100 个类别，每个类别有 600 张图像。其中 50,000 张用于训练，10,000 张用于测试。数据集的特点包括：

小尺寸图像：32×32 像素的图像尺寸较小，这使得模型需要更精细的特征提取能力。
细粒度分类：100 个类别中包含许多相似的子类（如不同种类的鱼类或花卉），增加了分类难度。
类别不均衡：虽然 CIFAR100 的类别分布相对均衡，但在实际应用中，自定义数据集可能会遇到严重的类别不均衡问题。

数据加载慢：由于图像尺寸小但数量多，数据加载和预处理可能成为瓶颈。
类别不均衡：某些类别的样本数量较少，可能导致模型偏向多数类。
小图像分类的模型适配：传统的卷积神经网络（CNN）可能需要对输入尺寸进行调整，以适配 32×32 的图像。

PyTorch 的 DataLoader 提供了多种参数配置来优化数据加载效率：

num_workers：设置多进程数据加载的进程数，建议设置为 CPU 核心数的 2 - 4 倍。
pin_memory：将数据加载到固定的内存区域，加速数据从 CPU 到 GPU 的传输。
batch_size：根据 GPU 内存选择合适的批量大小，通常从 64 或 128 开始尝试。

混合精度训练（Mixed Precision Training）通过使用 FP16 和 FP32 结合的方式，减少内存占用并加速计算。RAM 缓存则可以将部分数据预先加载到内存中，减少磁盘 I / O 的等待时间。

使用 torchvision.transforms 实现高效的数据增强策略，例如：

随机水平翻转
随机裁剪
颜色抖动
标准化

以下是一个完整的 ResNet-18 训练示例，包含自定义 Dataset 类实现、学习率 warmup 策略、类别权重采样和混合精度训练上下文管理。

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from torch.optim import Adam
from torch.cuda.amp import GradScaler, autocast

# 自定义 Dataset 类
class CIFAR100Dataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

# 数据加载和预处理
transform_train = transforms.Compose([transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
])

transform_test = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
])

# 加载数据集
trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)

# 数据加载器
trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)
testloader = DataLoader(testset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True)

# 初始化模型
model = torchvision.models.resnet18(pretrained=False)
model.fc = torch.nn.Linear(512, 100)
model = model.cuda()

# 优化器和学习率 warmup
optimizer = Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=len(trainloader), epochs=100)

# 混合精度训练
scaler = GradScaler()

# 训练循环
for epoch in range(100):
    model.train()
    for images, labels in trainloader:
        images, labels = images.cuda(), labels.cuda()
        optimizer.zero_grad()
        with autocast():
            outputs = model(images)
            loss = torch.nn.functional.cross_entropy(outputs, labels)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

# 测试
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        images, labels = images.cuda(), labels.cuda()
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')

优化前后 GPU 利用率对比：