How can I change part of a PyTorch tensor based on the values of another tensor?

user_239879 0
Viewed 1713

This question may not be clear, so please ask for clarification in the comments and I will expand.

I have the following tensors of the following shape:

mask.size() == torch.Size([1, 400])
clean_input_spectrogram.size() == torch.Size([1, 400, 161])
output.size() == torch.Size([1, 400, 161])

mask is comprised only of 0 and 1. Since it's a mask, I want to set the elements of output equal to clean_input_spectrogram where that relevant mask value is 1.

How would I do that?

pytorch python
2 Answers

You can do something like this, where:

  • m is your mask;
  • x is your spectogram;
  • o is your output;
import torch
torch.manual_seed(2020)

m = torch.tensor([[0, 1, 0]]).to(torch.int32)
x = torch.rand((1, 3, 2))
o = torch.rand((1, 3, 2))

print(o)
# tensor([[[0.5899, 0.8105],
#          [0.2512, 0.6307],
#          [0.5403, 0.8033]]])
print(x)
# tensor([[[0.4869, 0.1052],
#          [0.5883, 0.1161],
#          [0.4949, 0.2824]]])

o[:, m[0].to(torch.bool), :] = x[:, m[0].to(torch.bool), :]
# or 
# o[:, m[0] == 1, :] = x[:, m[0] == 1, :]

print(o)
# tensor([[[0.5899, 0.8105],
#          [0.5883, 0.1161],
#          [0.5403, 0.8033]]])
user_42282750

Basic usage of a mask:

a = torch.zeros([2,])     # tensor([0., 0.])
b = torch.ones([2,]) * 7  # tensor([7., 7.])

m = torch.tensor([True, False])

a[m] = b[m]               # tensor([7., 0.])

In this case, you have an extra dimension in a and b, but not in mask. This is no problem and will be handled automatically due to broadcasting. Example:

a = torch.zeros([2, 161])
b = torch.ones([2, 161]) * 7

m = torch.tensor([True, False])

a[m] = b[m]
print(a)
user_174385110
Related