PyTorch tensor indexing produces different results after same operation

username: Zhao_Liu
cooked: <p><strong>Short version:</strong></p>
<p>The same operation generates different result.</p>
<pre><code class="lang-auto">img1 = torch.zeros(3, img_size, img_size)
img2 = torch.zeros(3, img_size, img_size)
print(torch.allclose(img1, img2))
img1[:, pos[:, 1], pos[:, 0]] = color.T
img2[:, pos[:, 1], pos[:, 0]] = color.T
print('two imgs are identical:', torch.allclose(img1, img2))
----------output-----------
two imgs are identical: False
</code></pre>
<p><strong>Long version:</strong></p>
<p>Assume I have <code>n</code> particles in a space. Each particle has a position <code>(x, y, d)</code> where d prepresent the depth related to the camera and a color <code>(r, g, b)</code>. I am trying to render an image through these two tensors. So the shape of the tensors is <code>n*3</code> and <code>n*3</code>. For each pixel on the rendered image, the color should be the color of the shallowest particle with the same <code>(x, y)</code> coordinate. You can imagine the shallowest particle blocks the deeper ones.</p>
<p>Let’s create these two tensors</p>
<pre><code class="lang-auto">import torch
import torchvision
import time
import matplotlib.pyplot as plt

img_size = 50
num_particles = img_size*img_size*5
poses = torch.rand(num_particles, 3)*img_size
colors = torch.rand(num_particles, 3)
</code></pre>
<p>Assign colors of each particle according to the particle’s depth, the deeper the brighter. Thus, we can visualize the depth of the particle rendered on the image to check whether the code works or not.</p>
<pre><code class="lang-auto">for i in range(colors.shape[0]):
    colors[i, :] = poses[i, 2]/img_size
</code></pre>
<p>Sort the two tensors by depth. In this case, the output image should be close to black if descending=True, white otherwise.</p>
<pre><code class="lang-auto">_, indx = torch.sort(poses[:, 2], 0, descending=True)
pos = poses[indx].long()
color = colors[indx]
</code></pre>
<p>Now, I do the operation using a for loop to make sure my idea is correct.</p>
<pre><code class="lang-auto">%%time
img3 = torch.zeros(3, img_size, img_size)
for i in range(pos.shape[0]):
    img3[:, pos[i, 1], pos[i,0]] = color[i]
         
img4 = torch.zeros(3, img_size, img_size)
for i in range(pos.shape[0]):
    img4[:, pos[i, 1], pos[i,0]] = color[i]
---------ouptut------------
CPU times: user 269 ms, sys: 20.1 ms, total: 289 ms
Wall time: 271 ms
---------intput------------
print('two imgs are identical:', torch.allclose(img3, img4))
plt.imshow(torchvision.utils.make_grid([img3, img4]).permute(1, 2, 0))
plt.axis('off')
plt.show()
---------output-------------
two imgs are identical: True
</code></pre>
<p>However, using a for loop is very slow when the number of particles is large. Therefore, I want to verctorize the operation. I use <code>img[:, pos[:, 1], pos[:, 0]]</code> for the indexing. We can check this indexing should work in theory.</p>
<pre><code class="lang-auto">img = torch.rand(3, img_size, img_size)
a = img[:, pos[:, 1], pos[:, 0]]
time.sleep(1)
b = img[:, pos[:, 1], pos[:, 0]]
print('indexing works well:', torch.allclose(a, b))
print('shape matches:', img[:, pos[:, 1], pos[:, 0]].shape == color.T.shape)
c = color.T
d = color.T
print('transpose works well:', torch.allclose(c, d))
---------output----------
indexing works well: True
shape matches: True
transpose works well: True
</code></pre>
<p>But when I actually use this indexing operation to generate the image, the results are different.</p>
<pre><code class="lang-auto">%%time
img1 = torch.zeros(3, img_size, img_size)
img2 = torch.zeros(3, img_size, img_size)
img1[:, pos[:, 1], pos[:, 0]] = color.T
img2[:, pos[:, 1], pos[:, 0]] = color.T
---------output-------------
CPU times: user 34.9 ms, sys: 2.01 ms, total: 36.9 ms
Wall time: 3.09 ms
---------input--------------
print('two imgs are identical:', torch.allclose(img1, img2))
plt.imshow(torchvision.utils.make_grid([img1, img2]).permute(1, 2, 0))
plt.axis('off')
plt.show()
---------output-------------
two imgs are identical: False
</code></pre>
<p>What is wrong and how can I get the code work properly?</p>

username: ptrblck
cooked: <p>Could you check if you are trying to assign different values to the same index which might show the different results?</p>