python - 如何在 GPU 上从 PyTorch 中的其他两个张量有条件地构造一个张量？

Question

一个例子:

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 = tensor([1,2,1,1,0], device='cuda:0') .

基本上，当 pred与 correct 相同然后 correct否则 assigned .

此外，我想从梯度计算中排除这个计算。

有没有办法在不迭代张量的情况下做到这一点？

Answer 1

torch.where 完全符合您的要求:

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')

由于这些张量都没有 requires_grad=True ，不需要做任何事情来避免梯度计算。否则，您可以执行以下操作:

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')

如果您不使用 torch.no_grad() ，您将拥有:

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

然后，可以使用以下方法将其与计算图分离:

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