Torchvision example on Mixup and CrossEntropyLoss

Hongnan_G: <p>In PyTorch’s recent vision examples here: <a href="https://github.com/pytorch/vision/blob/main/references/classification/transforms.py" class="inline-onebox" rel="noopener nofollow ugc">vision/transforms.py at main · pytorch/vision · GitHub</a><br>
and <a href="https://github.com/pytorch/vision/blob/main/references/classification/train.py" class="inline-onebox" rel="noopener nofollow ugc">vision/train.py at main · pytorch/vision · GitHub</a>, it was shown how to use <code>Mixup</code> with the pipeline. I noticed when trying to use their <code>Mixup</code> function on my own that <code>CrossEntropyLoss</code> in general don’t expect targets to be of one-hot encoded, and it threw me a <code>RuntimeError: Expected object of scalar type Long but got scalar type Float for argument #2 'target' in call to _thnn_nll_loss_forward</code> error.</p>
<p>Am I missing something from the documentation</p>
<pre><code class="lang-python">class RandomMixup(torch.nn.Module):
    """Randomly apply Mixup to the provided batch and targets.
    The class implements the data augmentations as described in the paper
    `"mixup: Beyond Empirical Risk Minimization" &lt;https://arxiv.org/abs/1710.09412&gt;`_.
    Args:
        num_classes (int): number of classes used for one-hot encoding.
        p (float): probability of the batch being transformed. Default value is 0.5.
        alpha (float): hyperparameter of the Beta distribution used for mixup.
            Default value is 1.0.
        inplace (bool): boolean to make this transform inplace. Default set to False.
    """

    def __init__(self, num_classes: int, p: float = 0.5, alpha: float = 1.0, inplace: bool = False) -&gt; None:
        super().__init__()
        assert num_classes &gt; 0, "Please provide a valid positive value for the num_classes."
        assert alpha &gt; 0, "Alpha param can't be zero."

        self.num_classes = num_classes
        self.p = p
        self.alpha = alpha
        self.inplace = inplace

    def forward(self, batch: Tensor, target: Tensor) -&gt; Tuple[Tensor, Tensor]:
        """
        Args:
            batch (Tensor): Float tensor of size (B, C, H, W)
            target (Tensor): Integer tensor of size (B, )
        Returns:
            Tensor: Randomly transformed batch.
        """
        if batch.ndim != 4:
            raise ValueError(f"Batch ndim should be 4. Got {batch.ndim}")
        if target.ndim != 1:
            raise ValueError(f"Target ndim should be 1. Got {target.ndim}")
        if not batch.is_floating_point():
            raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.")
        if target.dtype != torch.int64:
            raise TypeError(f"Target dtype should be torch.int64. Got {target.dtype}")

        if not self.inplace:
            batch = batch.clone()
            target = target.clone()

        if target.ndim == 1:
            target = torch.nn.functional.one_hot(target, num_classes=self.num_classes).to(dtype=batch.dtype)

        if torch.rand(1).item() &gt;= self.p:
            return batch, target

        # It's faster to roll the batch by one instead of shuffling it to create image pairs
        batch_rolled = batch.roll(1, 0)
        target_rolled = target.roll(1, 0)

        # Implemented as on mixup paper, page 3.
        lambda_param = float(torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0])
        batch_rolled.mul_(1.0 - lambda_param)
        batch.mul_(lambda_param).add_(batch_rolled)

        target_rolled.mul_(1.0 - lambda_param)
        target.mul_(lambda_param).add_(target_rolled)

        return batch, target
</code></pre>
<p>where they did one hot encoded the targets.</p>

KFrank: <p>Hi Hongnan!</p>
<aside class="quote no-group" data-username="Hongnan_G" data-post="1" data-topic="142561" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hongnan_g/48/32949_2.png" class="avatar"> Hongnan_G:</div>
<blockquote>
<p><code>CrossEntropyLoss</code> in general don’t expect targets to be of one-hot encoded,</p>
</blockquote>
</aside>
<p>Possibly relevant to your issue:</p>
<p>As of <a href="https://pytorch.org/docs/1.10.0/generated/torch.nn.CrossEntropyLoss.html" rel="noopener nofollow ugc">version 1.10</a>, <code>CrossEntropyLoss</code> supports <em>probabilistic</em> targets (in<br>
addition to integer-class-label targets).  One-hot-encoded targets can be<br>
understood as a special case of probabilistic targets.  Earlier versions did<br>
not support probabilistic targets (and hence not one-hot-encoded targets).</p>
<p>If you are seeing your issue with an earlier version, you might try upgrading<br>
to the latest stable version, 1.10.1.</p>
<p>Best.</p>
<p>K. Frank</p>

Hongnan_G: <p>I think that is the answer, I am still using 1.9.1.</p>