Pytorch 中的图像增强

Question

我喜欢交替地增强图像。 我有pytorch转换代码如下。

import torchvision.transforms as tt
from torchvision.datasets import ImageFolder
#Data transform (normalization & data augmentation)
stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
train_tfms = tt.Compose([tt.RandomCrop(32, padding = 4, padding_mode = 'reflect'),
                         tt.RandomHorizontalFlip(),
                         tt.RandomAffine(degrees=(10, 30),
                                         translate=(0.1, 0.3),
                                         scale=(0.7, 1.3),
                                         shear=0.1, 
                                         resample=Image.BICUBIC)
                         tt.ToTensor(),
                         tt.Normalize(*stats)])

当我如下创建数据集并进行训练时，所有图像都会被增强。

train_ds = ImageFolder('content/train', train_tfms)

但我想要交替。第一张图像，作为原始图像进行训练。但是下一张图片被增强了。

我该怎么做？

Answer 1

您可以从单个数据集创建两个数据集，一个具有增强功能，另一个不具有增强功能，然后将它们连接起来。由于我们使用的是子数据集 pytorch 类，它将为我们处理这个顺序，因此将保持顺序。

train_ds_no_aug = ImageFolder('content/train')
train_ds_aug = ImageFolder('content/train', train_tfms)

# Check that aug_idx and no_aug_idx are not overlapping
aug_idx = torch.arange(1, len(train_ds_no_aug), 2)
no_aug_idx = torch.arange(0, len(train_ds_no_aug), 2)

train_ds_no_aug = torch.utils.data.Subset(train_ds_no_aug, no_aug_idx)
train_ds_aug = torch.utils.data.Subset(train_ds_aug, aug_idx)

train_ds = torch.utils.data.ChainDataset([train_ds_no_aug, train_ds_aug])
# Done :=