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

1次阅读

没有评论

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

CIFAR-10 是计算机视觉领域的经典基准数据集，包含以下特点：

图像规格 ：60,000 张 32×32 像素的彩色 RGB 图像
类别划分 ：10 个互斥类别（飞机、汽车、鸟、猫等），每个类别 6,000 张
标准拆分 ：50,000 张训练集 + 10,000 张测试集

该数据集常被用于：

图像分类模型的基准测试
轻量级网络架构验证
数据增强策略效果评估

实际使用中常遇到这些问题：

内存问题 ：
直接加载全部数据导致 OOM（尤其显存不足时）
未使用数据流式加载浪费内存
数据增强陷阱 ：
训练 / 验证集增强策略不一致
过度增强导致图像语义失真
类别不平衡 ：
某些类别样本量差异显著（实际场景常见）
模型偏向多数类预测

使用 PyTorch 的 DataLoader 配合自定义 Dataset：

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

# 定义转换管道
train_transform = transforms.Compose([transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))
])

# 创建数据集实例
train_set = datasets.CIFAR10(
    root='./data', 
    train=True,
    download=True, 
    transform=train_transform
)

# 使用 DataLoader 分批加载
train_loader = DataLoader(
    train_set, 
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=True  # 加速 GPU 传输
)

实现加权随机采样：

from torch.utils.data.sampler import WeightedRandomSampler

# 计算每个样本的权重
class_counts = [5000] * 10  # CIFAR-10 各类样本数
weights = 1. / torch.tensor(class_counts, dtype=torch.float)
samples_weights = weights[train_set.targets]

# 创建采样器
sampler = WeightedRandomSampler(
    weights=samples_weights,
    num_samples=len(samples_weights),
    replacement=True
)

# 修改 DataLoader 参数
train_loader = DataLoader(
    train_set,
    batch_size=128,
    sampler=sampler,  # 替换 shuffle
    num_workers=4
)

import torch.nn as nn

class BasicBlock(nn.Module):
    def __init__(self, in_planes, planes, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(
            in_planes, planes, kernel_size=3,
            stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3,
            stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, planes,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return F.relu(out)

from torch.optim import SGD
from torch.optim.lr_scheduler import CosineAnnealingLR

model = ResNet(BasicBlock, [2, 2, 2, 2]).cuda()
criterion = nn.CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = CosineAnnealingLR(optimizer, T_max=200)

for epoch in range(200):
    model.train()
    for inputs, targets in train_loader:
        inputs, targets = inputs.cuda(), targets.cuda()

        # 混合精度训练
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    scheduler.step()

    # 验证集评估
    model.eval()
    with torch.no_grad():
        correct = 0
        for inputs, targets in test_loader:
            outputs = model(inputs.cuda())
            pred = outputs.argmax(dim=1)
            correct += pred.eq(targets.cuda()).sum().item()

        acc = 100 * correct / len(test_set)
        print(f'Epoch {epoch}: Test Acc {acc:.2f}%')

通过 nvidia-smi 观察 GPU 利用率：

batch_size=64 → 约 40% 显存占用
batch_size=256 → 约 85% 显存占用
建议选择使 GPU 利用率达到 70-90% 的值

train_loader = DataLoader(
    train_set,
    batch_size=128,
    sampler=sampler,
    num_workers=4,
    prefetch_factor=2,  # 提前加载 2 个 batch
    persistent_workers=True
)

需安装 apex 库或使用 PyTorch 原生 AMP：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

错误做法 ：

# 错误：验证集不应使用数据增强
test_transform = transforms.Compose([transforms.RandomHorizontalFlip(),  # 不应该存在
    transforms.ToTensor(),
    transforms.Normalize(...)
])

正确做法 ：

test_transform = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize(...)  # 只做必要的标准化
])

训练集的 mean/std 应保存并用于验证集：

# 训练完成后保存参数
torch.save({'mean': [0.4914, 0.4822, 0.4465],
    'std': [0.247, 0.243, 0.261]
}, 'norm_params.pth')

# 验证时加载
params = torch.load('norm_params.pth')
test_transform = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize(params['mean'], params['std'])
])

可视化检查数据增强效果：

import matplotlib.pyplot as plt

def imshow(img):
    img = img * 0.247 + 0.4914  # 反标准化
    plt.imshow(img.permute(1, 2, 0))
    plt.show()

# 检查第一个 batch
images, _ = next(iter(train_loader))
imshow(images[0])