from abc import ABCMeta, abstractmethod import torch import torch.nn.functional as F from torch.distributions.beta import Beta class BaseMixupLayer(object, metaclass=ABCMeta): """Base class for MixupLayer""" def __init__(self): super(BaseMixupLayer, self).__init__() @abstractmethod def mixup(self, imgs, gt_label): pass class BatchMixupLayer(BaseMixupLayer): """Mixup layer for batch mixup. Args: alpha (float): Parameters for Beta distribution. num_classes (int): The number of classes. """ def __init__(self, alpha, num_classes): super(BatchMixupLayer, self).__init__() assert isinstance(alpha, float) assert isinstance(num_classes, int) self.alpha = alpha self.num_classes = num_classes self.Beta = Beta(self.alpha, self.alpha) def mixup(self, img, gt_label): lam = self.Beta.sample() batch_size = img.size(0) index = torch.randperm(batch_size) one_hot_gt_label = F.one_hot(gt_label, num_classes=self.num_classes) mixed_img = lam * img + (1 - lam) * img[index, :] mixed_gt_label = lam * one_hot_gt_label + ( 1 - lam) * one_hot_gt_label[index, :] return mixed_img, mixed_gt_label def __call__(self, img, gt_label): return self.mixup(img, gt_label)