7

This is a follow up question to this question. I want to do the exactly same thing in pytorch. Is it possible to do this? If yes, how?

import torch
image = torch.tensor([[246,  50, 101], [116,   1, 113], [187, 110,  64]])
iy = torch.tensor([[1, 0, 2], [1, 0, 2], [2, 2, 2]])
ix = torch.tensor([[0, 2, 1], [1, 2, 0], [0, 1, 2]])
warped_image = torch.zeros(size=image.shape)

I need something like torch.add.at(warped_image, (iy, ix), image) that gives the output as

[[  0.   0.  51.]
 [246. 116.   0.]
 [300. 211.  64.]]

Note that the indices at (0,1) and (1,1) point to the same location (0,2). So, I want warped_image[0,2] = image[0,1] + image[1,1] = 51.

Nagabhushan S N
  • 6,407
  • 8
  • 44
  • 87

1 Answers1

6

What you are looking for is torch.Tensor.index_put_ with the accumulate argument set to True:

>>> warped_image = torch.zeros_like(image)

>>> warped_image.index_put_((iy, ix), image, accumulate=True)
tensor([[  0,   0,  51],
        [246, 116,   0],
        [300, 211,  64]])

Or, using the out-place version torch.index_put:

>>> torch.index_put(torch.zeros_like(image), (iy, ix), image, accumulate=True)
tensor([[  0,   0,  51],
        [246, 116,   0],
        [300, 211,  64]])
Ivan
  • 34,531
  • 8
  • 55
  • 100
  • Wow! Thank you so much! I hope this allows the gradients to flow. – Nagabhushan S N Jan 05 '21 at 18:32
  • 1
    Yes! If you set `requires_grad` to `True` on `image`, then `warped_image` will turn out with a `grad_fn` (used to backpropagate) set to ``. – Ivan Jan 05 '21 at 18:37
  • I need the gradients to flow through the indices as well, but I'm getting an error for that. Any idea how to solve it or any workarounds? I've asked a separate [question](https://stackoverflow.com/q/66973502/3337089) – Nagabhushan S N Apr 06 '21 at 17:35
  • 1
    @NagabhushanSN I don't think you can pass a gradient through the indices, as those are discrete and non-differentiable – Ryan Burgert Jan 06 '22 at 01:29
  • Yeah, I figured that. And it turned out that that wasn't needed after all. – Nagabhushan S N Jan 06 '22 at 05:01