如何使用 CutMix 和 MixUp#

CutMixMixUp 是流行的增强策略,可以提高分类精度。

这些变换与Torchvision的其他变换略有不同,因为它们期望输入为样本批次,而不是单个图像。在这个例子中,我们将解释如何使用它们:在 DataLoader 之后,或作为整理函数的一部分。

import torch
from torchvision.datasets import FakeData
from torchvision.transforms import v2


NUM_CLASSES = 100

预处理流程#

将使用简单但典型的图像分类处理流程:

preproc = v2.Compose([
    v2.PILToTensor(),
    v2.RandomResizedCrop(size=(224, 224), antialias=True),
    v2.RandomHorizontalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=True),  # to float32 in [0, 1]
    v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),  # typically from ImageNet
])

dataset = FakeData(size=1000, num_classes=NUM_CLASSES, transform=preproc)

img, label = dataset[0]
print(f"{type(img) = }, {img.dtype = }, {img.shape = }, {label = }")
type(img) = <class 'torch.Tensor'>, img.dtype = torch.float32, img.shape = torch.Size([3, 224, 224]), label = 67

值得注意的是,CutMix 和 MixUp 都不是这个预处理管道的一部分。一旦定义了DataLoader,我们将稍后加入它们。为了复习一下,如果不使用 CutMix 或 MixUp,DataLoader 和训练循环看起来会是这样:

from torch.utils.data import DataLoader

dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

for images, labels in dataloader:
    print(f"{images.shape = }, {labels.shape = }")
    print(labels.dtype)
    # <rest of the training loop here>
    break
images.shape = torch.Size([4, 3, 224, 224]), labels.shape = torch.Size([4])
torch.int64

MixUp和CutMix的应用场景#

在DataLoader之后使用MixUp和CutMix#

现在让我们添加CutMix和MixUp。最简单的方法是在DataLoader之后立即这样做:Dataloader已经为我们批处理了图像和标签,这正是这些转换期望的输入格式:

dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

cutmix = v2.CutMix(num_classes=NUM_CLASSES)
mixup = v2.MixUp(num_classes=NUM_CLASSES)
cutmix_or_mixup = v2.RandomChoice([cutmix, mixup])

for images, labels in dataloader:
    print(f"Before CutMix/MixUp: {images.shape = }, {labels.shape = }")
    images, labels = cutmix_or_mixup(images, labels)
    print(f"After CutMix/MixUp: {images.shape = }, {labels.shape = }")

    # <rest of the training loop here>
    break
Before CutMix/MixUp: images.shape = torch.Size([4, 3, 224, 224]), labels.shape = torch.Size([4])
After CutMix/MixUp: images.shape = torch.Size([4, 3, 224, 224]), labels.shape = torch.Size([4, 100])

请注意标签是如何被转换的:我们从批处理标签的形状(batch_size,)转变为形状为(batch_size, num_classes)的张量。这些经过转换的标签仍然可以直接传递给损失函数,例如 torch.nn.functional.cross_entropy()

作为集合函数的一部分#

在DataLoader之后传递转换操作是使用CutMix和MixUp的最简单方式,但一个缺点是它没有利用DataLoader的多进程优势。为此,我们可以将这些转换作为合并函数的一部分传递(请参阅PyTorch文档以了解更多关于合并的信息)。

from torch.utils.data import default_collate


def collate_fn(batch):
    return cutmix_or_mixup(*default_collate(batch))


dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2, collate_fn=collate_fn)

for images, labels in dataloader:
    print(f"{images.shape = }, {labels.shape = }")
    # No need to call cutmix_or_mixup, it's already been called as part of the DataLoader!
    # <rest of the training loop here>
    break
images.shape = torch.Size([4, 3, 224, 224]), labels.shape = torch.Size([4, 100])

非标准输入格式#

迄今为止,我们使用了典型的样本结构,其中输入为(images, labels)。MixUp和CutMix默认情况下可以神奇地与大多数常见的样本结构一起工作:第二个参数是张量标签的元组,或者是包含"label[s]"键的字典。有关更多详细信息,请查看labels_getter参数的文档。

如果您的样本具有不同的结构,您仍然可以通过向labels_getter参数传递一个可调用对象来使用CutMix和MixUp。例如:

batch = {
    "imgs": torch.rand(4, 3, 224, 224),
    "target": {
        "classes": torch.randint(0, NUM_CLASSES, size=(4,)),
        "some_other_key": "this is going to be passed-through"
    }
}


def labels_getter(batch):
    return batch["target"]["classes"]


out = v2.CutMix(num_classes=NUM_CLASSES, labels_getter=labels_getter)(batch)
print(f"{out['imgs'].shape = }, {out['target']['classes'].shape = }")
out['imgs'].shape = torch.Size([4, 3, 224, 224]), out['target']['classes'].shape = torch.Size([4, 100])