PyTorch: Perform add/sub/mul operations using 1-D tensor and multi-channel (3-D) image tensor

Question

Note: I am looking for the fastest/optimal way of doing or improving on this, as my constraint is time

Using PyTorch, I have an image as a 3D tensor, lets say of dimension 64 x 400 x 400, where 64 refers to channels and 400 is the image dimensions. With this, I have a 64 length 1D tensor, all with different values, where the intention is to use one value per channel. I want to use a value per 1D tensor to apply to the entire 400x400 block of the channel. So for example, when I want to add 3d_tensor + 1d_tensor, I want 1d_tensor[i] to be added to all 400x400 = 160000 values in 3d_tensor[i], with [i] ranging from 0 to 63.

What I previously did:
I tried doing it directly, by using the operators only:

output_add = 1d_tensor + 3d_tensor

This returned an error that the dimension of 3d_tensor (400) and 1d_tensor (64) are incompatible.

So my current form is using a for loop

for a, b in zip(3d_tensor, 1d_tensor):
   a += b

However at one stage I have four different 1D tensors to use at once in either addition, subtraction or multiplication, so is this for loop method the most efficient? I'm also planning on doing it 20+ times per image, so speed is key. I tried maybe extending the 1d tensor to a dimension of 64 x 400 x 400 also, so it could be used directly, but I could not get this right using tensor.repeat()

Answer 1

You should add some dimension to the 1D array: convert from 64 to (64 x 1 x 1):

output_add = 1d_tensor[:, None, None] + 3d_tensor

with this None type indexing you can add dimension any where. The [:, None, None] will add two additional dimension to the 1D array after the existing dimension.

Or you can use view for the same result see:

output_add = 1d_tensor.view(-1, 1, 1) + 3d_tensor

The result has the same dimension as the 3D array: (64, 1, 1). So the pytorch can use broadcasting.

Here a good explanation for broad casting: How does pytorch broadcasting work?