How to replace scipy.ndimage.rotate with torchvision

sakura1: <p>I am training a vision model with  scipy.ndimage.rotate in cpu-aug. This rotation in the cpu becomes the bottleneck. So I was looking for suggestions to replace it with torchvision. But the resut is diffent from torchvision.transforms.functional.rotate.<br>
My test looks something like this:</p>
<pre><code class="lang-auto">import numpy as np
from scipy import ndimage

c = np.zeros([4, 4]).astype("float32")
for i in range(4):
    for j in range(4):
        c[i][j] = 4 * i + j

res1 = ndimage.rotate(c, 60, reshape=False, mode="nearest")
print(res1)

import torchvision
from torchvision.transforms import InterpolationMode
inp = torch.tensor(c, device="cuda:0")
res2 = torchvision.transforms.functional.rotate(inp[None][None], 60, interpolation=InterpolationMode.NEAREST)

print(res2)
</code></pre>
<p>The result  looks something like this:</p>
<pre><code class="lang-auto">res1 :
[[ 2.0601943  3.3288476  7.8078837 11.240777 ]
 [ 1.2019709  4.727933   8.9034605 13.673849 ]
 [ 1.3261505  6.096539  10.272067  13.798029 ]
 [ 3.759223   7.1921163 11.671152  12.939806 ]]

res2:
tensor([[[[ 0.,  3.,  7.,  0.],
          [ 1.,  6., 10., 15.],
          [ 0.,  5.,  9., 14.],
          [ 0.,  8., 12.,  0.]]]], device='cuda:0')

</code></pre>
<p>Thanks for taking the time to read through this! Any suggestion is welcome</p>

ptrblck: <p>The <code>ndimage</code> result seems unexpected to me. It seems you are creating an example input array using integer values represented in <code>float32</code>. Both transformation use the <code>nearest</code> interpolation mode which should pick the nearest input value for the corresponding output. Since all inputs are integers I would expect to also see integers only in the transformed output. While this is the case for <code>torchvision</code> the <code>ndimage</code> transformation returns floating point values.<br>
I’m currently not in front of my workstation so cannot reproduce or debug it.</p>
<p>EDIT: I was wrong and the <code>mode</code> argument in <code>ndimage</code> defines the padding behavior. From the docs:</p>
<pre><code class="lang-python">The mode parameter determines how the input array is extended beyond its boundaries. Default is ‘constant’. Behavior for each valid value is as follows (see additional plots and details on boundary modes)
</code></pre>
<p>and it seems a spline interpolation is used by default.</p>

sakura1: <p>Thank you for answer.  If input array represented in <code>uint8</code>， the result looks different：</p>
<pre><code class="lang-auto">c = np.zeros([4, 4]).astype("uint8")

res1:
[[ 2  3  8 11]
 [ 1  5  9 14]
 [ 1  6 10 14]
 [ 4  7 12 13]]

res2:
tensor([[[[ 0,  3,  7,  0],
          [ 1,  6, 10, 15],
          [ 0,  5,  9, 14],
          [ 0,  8, 12,  0]]]], device='cuda:0', dtype=torch.uint8)
</code></pre>
<p>Both interpolation result is different, although both are nearest.</p>
<p>My question:<br>
1、 How can i repacle scipy.ndimage.rotate with torchvision operations (affine？)</p>

ptrblck: <aside class="quote no-group" data-username="sakura1" data-post="3" data-topic="182342">
<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/sakura1/48/74707_2.png" class="avatar"> sakura1:</div>
<blockquote>
<p>Both interpolation result is different, although both are nearest.</p>
</blockquote>
</aside>
<p>They are not as explained in my previous post, as <code>ndimage.rotate</code> uses a spline interpolation.<br>
You could disable it and experiment further to narrow down if more settings need to be changed:</p>
<pre data-code-wrap="python"><code class="lang-python">res1 = ndimage.rotate(c, 60, reshape=False, mode="constant", cval=0, order=0, prefilter=False)
</code></pre>
<p>This setup reduces the error, but still shows 2 mismatches as the edges:</p>
<pre data-code-wrap="python"><code class="lang-python">[[ 0  3  0  0]
 [ 0  6 10 15]
 [ 0  5  9  0]
 [ 0  0 12  0]]
tensor([[[[ 0,  3,  7,  0],
          [ 1,  6, 10, 15],
          [ 0,  5,  9, 14],
          [ 0,  8, 12,  0]]]], device='cuda:0', dtype=torch.uint8)
</code></pre>