torch.where 的详细用法

 

torch.where 会根据条件去选择元素，返回一个tensor。¶

torch.where(condition, input, other, *, out=None) → Tensor

condition 是一个BoolTensor

input 和 other 可以常量，也可以是张量,

返回的张量形状与condition相同

 

In [2]:

import torch
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"

a = torch.randint(0,10,[2,3])
torch.where(a > 5, 6,7)

 

Out[2]:

tensor([[6, 7, 6],
        [6, 7, 7]])

 

 

当input为张量的情况

 

In [9]:

a = torch.randint(0,10,[2,3])
a
b = torch.rand([1,3])  # b张量会被广播成(2,3)的形状
b
torch.where(a > 5, b,7)

 

Out[9]:

tensor([[7, 4, 1],
        [1, 6, 3]])

Out[9]:

tensor([[0.3639, 0.3413, 0.0434]])

Out[9]:

tensor([[0.3639, 7.0000, 7.0000],
        [7.0000, 0.3413, 7.0000]])

 

 

当input 和 other 都为张量的时候

 

In [11]:

a = torch.randint(0,10,[2,3])
a
b = torch.rand([1,3])  # b张量会被广播成(2,3)的形状
b
c = torch.randint(10,20,[1,3])
c
torch.where(a > 5, b,c)

 

Out[11]:

tensor([[7, 1, 5],
        [1, 7, 4]])

Out[11]:

tensor([[0.4111, 0.3227, 0.3590]])

Out[11]:

tensor([[19, 16, 19]])

Out[11]:

tensor([[ 0.4111, 16.0000, 19.0000],
        [19.0000,  0.3227, 19.0000]])

 

 

当不传入 input 和 other¶

torch.where(condition) 会返回满足 condition 为非零元素的坐标.

等同于 torch.nonzero

 

In [26]:

a = torch.randint(0,5,[2,3])
a
b,c = torch.where(a) # b 和 c 的长度相等，表示非零元素的数量。
b  #行索引
c  #列索引
for b_index in range(len(b)):
    a[b[b_index], c[b_index]]
    

 

Out[26]:

tensor([[4, 4, 3],
        [0, 1, 0]])

Out[26]:

tensor([0, 0, 0, 1])

Out[26]:

tensor([0, 1, 2, 1])

Out[26]:

tensor(4)

Out[26]:

tensor(4)

Out[26]:

tensor(3)

Out[26]:

tensor(1)