尽管看上去 MixUp 和 Mosaic 只是一种特殊的数据变换,按理说应该在 transform 里实现。考虑到这两种增强会涉及到“将多张图片融合成一张图片”的操作,在 transform 里实现他们的难度就会很大,因为目前 transform 的范式是对一张图片做各种增强,我们很难在一个 transform 里去额外读取其他图片(transform 里无法访问到 dataset)。
class MixUpDataPreprocessor(nn.Module):
def __init__(self, num_class, alpha):
self.alpha = alpha
def forward(self, data, training=True):
data = tuple(_data.cuda() for _data in data)
# 验证阶段无需进行 MixUp 数据增强
if not training:
return data
label = F.one_hot(label) # label 转 onehot 编码
batch_size = len(label)
index = torch.randperm(batch_size) # 生成一个0-batch_size的随机排列
img, label = data
lam = np.random.beta(self.alpha, self.alpha) # 融合因子
# 原图和标签的 MixUp.
img = lam * img + (1 - lam) * img[index, :]
label = lam * batch_scores + (1 - lam) * batch_scores[index, :]
# 由于此时返回的是 onehot 编码的 label,model 的 forward 也需要做相应调整
return tuple(img, label)