如何在PyTorch中对CIFAR-10部分图像做数据增强并绘制对比图

核心思路 #

不要在数据集初始化时传入全局增强transform，先加载无增强的原始数据集，按需抽取指定数量的样本后单独对其应用增强操作，再做对比可视化即可。

完整实现步骤 #

1. 导入依赖包 #

import torch
import matplotlib.pyplot as plt
from torchvision import datasets, transforms

2. 加载无增强的原始CIFAR-10数据集 #

这一步只做基础的格式转换，不添加任何增强操作，也可以选择暂不做归一化，方便后续可视化：

# 基础transform：仅转成Tensor，不做增强和归一化
basic_transform = transforms.Compose([
    transforms.ToTensor()
])
# 加载原始训练集
original_train_set = datasets.CIFAR10(
    root='./data/',
    train=True,
    download=True,
    transform=basic_transform
)

3. 定义你需要使用的数据增强规则 #

可以根据需求添加旋转、裁剪、翻转等操作：

# 自定义增强transform，这里加上你需要的随机水平翻转、随机旋转
augment_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.8), # 调高概率方便看到增强效果
    transforms.RandomRotation(degrees=15),  # 随机旋转±15度
    # 如果需要归一化可以加在这里，可视化的时候记得反归一化即可
    # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

4. 抽取指定数量的样本，单独做增强 #

比如我们取前10张图像做增强：

sample_num = 10
# 抽取前sample_num个样本，也可以自定义索引列表取指定样本
original_samples = [original_train_set[i][0] for i in range(sample_num)]
# 对每个样本单独应用增强
augmented_samples = [augment_transform(img) for img in original_samples]

5. 绘制原始图像与增强图像的对比图 #

plt.figure(figsize=(20, 6))
# 第一行：原始图像
for i in range(sample_num):
    plt.subplot(2, sample_num, i+1)
    # 把CHW格式的Tensor转成HWC的numpy格式用于画图
    plt.imshow(original_samples[i].permute(1, 2, 0).numpy())
    plt.title(f"Original {i+1}")
    plt.axis('off')
# 第二行：增强后的图像
for i in range(sample_num):
    plt.subplot(2, sample_num, sample_num + i + 1)
    plt.imshow(augmented_samples[i].permute(1, 2, 0).numpy())
    plt.title(f"Augmented {i+1}")
    plt.axis('off')
plt.tight_layout()
plt.show()

注意事项 #

• 如果你的增强流程里加了Normalize操作，可视化的时候需要先做反归一化，把数值拉回[0,1]区间，不然图像颜色会异常：

  # 反归一化函数示例，对应常用的ImageNet均值方差
  def denormalize(img_tensor):
      mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
      std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
      return img_tensor * std + mean
  

  画图的时候调用denormalize(augmented_samples[i])再转格式即可。
• 你之前代码里写的transform=data_transforms['train']，该写法前提是data_transforms是包含train/val键的字典，如果你只定义了单个变换pipeline直接传transform=data_transforms即可。

内容的提问来源于stack exchange，提问作者Raha Moosavi