How to conditionally construct a tensor from two other tensors in PyTorch on the GPU?

user_6392022 0
Viewed 244

An example:

import torch

pred = torch.tensor([1,2,1,0,0], device='cuda:0')
correct = torch.tensor([1,0,1,1,0], device='cuda:0')
assigned = torch.tensor([1,2,2,1,0], device='cuda:0')

I want result = tensor([1,2,1,1,0], device='cuda:0').

  1. Basically, when pred is same as correct then correct else assigned.

  2. Moreover, I want to exclude this computation from the gradient computation.

Is there a way to do this without iterating over the tensors?

pytorch python
1 Answers

torch.where does exactly what you're looking for:

import torch

pred = torch.tensor([1,2,1,0,0], device='cuda:0')
correct = torch.tensor([1,0,1,1,0], device='cuda:0')
assigned = torch.tensor([1,2,2,1,0], device='cuda:0')

result = torch.where(pred == correct, correct, assigned)

print(result)
# >>> tensor([1, 2, 1, 1, 0], device='cuda:0')

Since none of these tensors have requires_grad=True, nothing needs to be done to avoid gradient computation. Otherwise, you can do something like this:

import torch

pred = torch.tensor([1.,2.,1.,0.,0.], device='cuda:0')
correct = torch.tensor([1.,0.,1.,1.,0.], device='cuda:0', requires_grad=True)
assigned = torch.tensor([1.,2.,2.,1.,0.], device='cuda:0', requires_grad=True)

with torch.no_grad():
    result = torch.where(pred == correct, correct, assigned)

print(result)
# >>> tensor([1, 2, 1, 1, 0], device='cuda:0')

If you do not use torch.no_grad(), you'll have:

result = torch.where(pred == correct, correct, assigned)
print(result)
# >>> tensor([1., 2., 1., 1., 0.], device='cuda:0', grad_fn=<SWhereBackward>)

which, then, can be detached from the computational graph by using:

result = result.detach()
print(result)
# >>> tensor([1., 2., 1., 1., 0.], device='cuda:0')
user_42282750
Related