How to augment image tensor directly?

KeisukeShimokawa: <p>I am currently working on an implementation of the  Consistency regularization method proposed in the following paper.</p>
<p><a href="https://arxiv.org/abs/2002.04724" rel="noopener nofollow ugc">[arXiv:2002.04724] Improved Consistency Regularization for GAN</a></p>
<p>Within this paper, when injecting real and generated images into the Discriminator, they apply data augmentation to each image and compute the loss value according to the following algorithm.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/6/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f_2_690x463.jpeg" alt="image" data-base62-sha1="ci5Jxep4kMxhQ5aat8xqpapPHDF" width="690" height="463" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f_2_690x463.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f_2_1035x694.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/6/5625a743dd86cf0545aa8e8188b92a3b0ee0e53f.jpeg 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1065×715 206 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>If I only want to perform data augmentation on the real images, I can do so during batch data creation in <code>torch.utils.data.Dataset</code>, but how can I perform data augmentation on a Tensor generated from the Generator?</p>
<p>The following methods involve sending data between the CPU and the GPU, which causes some overhead and reduces computational speed.</p>
<pre data-code-wrap="python"><code class="lang-python">def transforms_images(image_tensor):
    """transform tensor for consistency regularization
    :param image_tensor: torch.Tensor [B, C, H, W]
    """

    device = image_tensor.device
    _, _, H, W = image_tensor.shape

    cr_transforms = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomResizedCrop((H, W)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, ), (0.5, ))
    ])

    image_tensor = torch.cat([cr_transforms(image).unsqueeze_(0) 
                              for image in image_tensor.cpu()], dim=0)

    return image_tensor.to(device=device)
</code></pre>
<p>What’s the best way?</p>

Nikronic: <p>Hi,</p>
<p>Obviously, as you have mentioned, transferring every tensor between cpu and gpu to use available methods which only works for PIL images are is proper.<br>
As far as I know, there is no built-in function for <code>Crop</code> but for others we have a solution.</p>
<p>One solution is to copy the source code and just change the PIL input to tensor. For instance, here is the implementation of <a href="https://pytorch.org/docs/stable/_modules/torchvision/transforms/transforms.html#RandomResizedCrop" rel="nofollow noopener">RandomResizedCrop</a><br>
For cropping, indexing just works fine and for resizing there is a built in function. The only issue is random generation for crop which can be copy pasted from source code.</p>
<p>Here is what I have changed that works for a single 3D tensor:</p>
<pre><code class="lang-auto">
# a arbitrary 3D input
x = torch.ones((3, 100, 100))*255
x[:, 25:75, 25:75] = 0

scale=(0.08, 1.0) 
ratio=(3. / 4., 4. / 3.)
width, height = x.shape[-2], x.shape[-1]
size = (64, 64)


area = height * width

for _ in range(10):
    target_area = random.uniform(*scale) * area
    log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
    aspect_ratio = math.exp(random.uniform(*log_ratio))
    z = None
    w = int(round(math.sqrt(target_area * aspect_ratio)))
    h = int(round(math.sqrt(target_area / aspect_ratio)))

    if 0 &lt; w &lt;= width and 0 &lt; h &lt;= height:
        i = random.randint(0, height - h)
        j = random.randint(0, width - w)
        z = i, j, h, w
        break
if z is None:
    # Fallback to central crop
    in_ratio = float(width) / float(height)
    if (in_ratio &lt; min(ratio)):
        w = width
        h = int(round(w / min(ratio)))
    elif (in_ratio &gt; max(ratio)):
        h = height
        w = int(round(h * max(ratio)))
    else:  # whole image
        w = width
        h = height
    i = (height - h) // 2
    j = (width - w) // 2
    z = i, j, h, w
resized = F.interpolate(x[:, i:h, j:w].unsqueeze(0), size=size, mode='bicubic')[0]
</code></pre>
<p>Also you can wrap your code in the form of source code I referenced.</p>
<p>Bests</p>

KeisukeShimokawa: <p>Thanks!</p>
<p>I tested it the way you presented it here.<br>
and fix some minor changes.</p>
<pre><code class="lang-python"># x[:, i:h, j:w]
x[:, i:i+h, j:j+w]
</code></pre>
<p><a href="https://colab.research.google.com/github/KeisukeShimokawa/papers-challenge/blob/master/src/gan/CRGAN/notebooks/RandomResizedCrop_Tensor.ipynb" rel="nofollow noopener">[Colab] RandomResizedCrop_Tensor.ipynb</a></p>
<p>it works!</p>

Nikronic: <p>That is probably making more sense as <code>i+h</code> or <code>j+w</code> never exceeds original <code>height</code> or <code>width</code> and those are only a proportion of original sizes.<br>
Although the code I provided needs few other fixes, for instance it uses external libraries that need to be replaced by <code>torch</code> for working on tensors.</p>
<p>Good luck</p>

Nikronic: <p>Hi, I just found a much better solution using <a href="https://kornia.readthedocs.io/en/latest/" rel="nofollow noopener">kornia</a>, the augmentation library that supports direct tensor and works on cuda too.</p>
<p>Here is the simplified solution:</p>
<pre><code class="lang-auto">import torch
import kornia as K

x = torch.randn(1, 3, 200, 200).cuda()
K.augmentation.RandomResizedCrop((65, 65))(x)
</code></pre>
<p>Note that this augmentation methods can be used in the exact same way as <code>nn.Module</code> within <code>nn.Sequential</code> layer.</p>
<p>Bests</p>

KeisukeShimokawa: <p>Very Good!!!<br>
Thanks!!!</p>