Short version:
The same operation generates different result.
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
Long version:
Assume I have n particles in a space. Each particle has a position (x, y, d) where d prepresent the depth related to the camera and a color (r, g, b). I am trying to render an image through these two tensors. So the shape of the tensors is n*3 and n*3. For each pixel on the rendered image, the color should be the color of the shallowest particle with the same (x, y) coordinate. You can imagine the shallowest particle blocks the deeper ones.
Let’s create these two tensors
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)
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.
for i in range(colors.shape[0]):
colors[i, :] = poses[i, 2]/img_size
Sort the two tensors by depth. In this case, the output image should be close to black if descending=True, white otherwise.
_, indx = torch.sort(poses[:, 2], 0, descending=True)
pos = poses[indx].long()
color = colors[indx]
Now, I do the operation using a for loop to make sure my idea is correct.
%%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
However, using a for loop is very slow when the number of particles is large. Therefore, I want to verctorize the operation. I use img[:, pos[:, 1], pos[:, 0]] for the indexing. We can check this indexing should work in theory.
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
But when I actually use this indexing operation to generate the image, the results are different.
%%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
What is wrong and how can I get the code work properly?