Class CutMix

monai/transforms/regularization/array.py:102–166

CutMix augmentation as described in: Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo. CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features, ICCV 2019 Class derived from :py:class:`monai.transforms.M

class CutMix(Mixer):
    """CutMix augmentation as described in:
        Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo.
        CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features,
        ICCV 2019

        Class derived from :py:class:`monai.transforms.Mixer`. See corresponding
        documentation for details on the constructor parameters. Here, alpha not only determines
        the mixing weight but also the size of the random rectangles used during for mixing.
        Please refer to the paper for details.

        Please note that there is a change in behavior starting from version 1.4.0. In the previous
        implementation, the transform would generate a different label each time it was called.
        To ensure determinism, the new implementation will now generate the same label for
        the same input image when using the same operation.

        The most common use case is something close to:

    .. code-block:: python

        cm = CutMix(batch_size=8, alpha=0.5)
        for batch in loader:
            images, labels = batch
            augimg, auglabels = cm(images, labels)
            output = model(augimg)
            loss = loss_function(output, auglabels)
            ...

    """

    def apply(self, data: torch.Tensor):
        weights, perm, coords = self._params
        nsamples, _, *dims = data.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

        return mask * data + (1 - mask) * data[perm, ...]

    def apply_on_labels(self, labels: torch.Tensor):
        weights, perm, _ = self._params
        nsamples, *dims = labels.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mixweight = weights[(Ellipsis,) + (None,) * len(dims)]
        return mixweight * labels + (1 - mixweight) * labels[perm, ...]

    def __call__(self, data: torch.Tensor, labels: torch.Tensor | None = None, randomize=True):
        data_t = convert_to_tensor(data, track_meta=get_track_meta())
        augmented_label = None
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())