如何使用 CutMix 和 MixUp#
CutMix
和 MixUp
是流行的增强策略,可以提高分类精度。
这些变换与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])