使用torch.where函數(shù)

首先我們看一下Pytorch中torch.where函數(shù)是怎樣定義的：

@overload
def where(condition: Tensor) -> Union[Tuple[Tensor, ...], List[Tensor]]: ...

torch.where函數(shù)的功能如下：

torch.where(condition, x, y)

• condition：判斷條件
• x：若滿足條件，則取x中元素
• y：若不滿足條件，則取y中元素

以具體實例看一下torch.where函數(shù)的效果：

import torch
 
# 條件
condition = torch.rand(3, 2)
print(condition)
# 滿足條件則取x中對應元素
x = torch.ones(3, 2)
print(x)
# 不滿足條件則取y中對應元素
y = torch.zeros(3, 2)
print(y)
# 條件判斷后的結果
result = torch.where(condition > 0.5, x, y)
print(result)

結果如下：

tensor([[0.3224, 0.5789],
        [0.8341, 0.1673],
        [0.1668, 0.4933]])
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]])
tensor([[0., 0.],
        [0., 0.],
        [0., 0.]])
tensor([[0., 1.],
        [1., 0.],
        [0., 0.]])

可以看到torch.where函數(shù)會對condition中的元素逐一進行判斷，根據(jù)判斷的結果選取x或y中的值，所以要求x和y應該與condition形狀相同。

torch.where(),np.where()兩種用法，及np.argwhere()尋找張量（tensor）和數(shù)組中為0的索引

1.torch.where()

torch.where()有兩種用法，

• 當輸入?yún)?shù)為三個時，即torch.where(condition, x, y)，返回滿足 x if condition else y的tensor，注意x,y必須為tensor
• 當輸入?yún)?shù)為一個時，即torch.where(condition)，返回滿足condition的tensor索引的元組(tuple)

代碼示例

torch.where(condition, x, y)

代碼

import torch
import numpy as np
 
# 初始化兩個tensor
x = torch.tensor([
    [1,2,3,0,6],
    [4,6,2,1,0],
    [4,3,0,1,1]
])
y = torch.tensor([
    [0,5,1,4,2],
    [5,7,1,2,9],
    [1,3,5,6,6]
])
 
# 尋找滿足x中大于3的元素，否則得到y(tǒng)對應位置的元素
arr0 = torch.where(x>=3, x, y) #輸入?yún)?shù)為3個
 
print(x, '\n', y)
print(arr0, '\n', type(arr0))

結果

>>> x
tensor([[1, 2, 3, 0, 6],
        [4, 6, 2, 1, 0],
        [4, 3, 0, 1, 1]])
>>> y
tensor([[0, 5, 1, 4, 2],
        [5, 7, 1, 2, 9],
        [1, 3, 5, 6, 6]])
 
>>> arr0
tensor([[0, 5, 3, 4, 6],
        [4, 6, 1, 2, 9],
        [4, 3, 5, 6, 6]])
 
>>> type(arr0)
<class 'torch.Tensor'>

arr0的類型為<class 'torch.Tensor'>

torch.where(condition)

以尋找tensor中為0的索引為例

代碼

import torch
import numpy as np
x = torch.tensor([
    [1,2,3,0,6],
    [4,6,2,1,0],
    [4,3,0,1,1]
])
y = torch.tensor([
    [0,5,1,4,2],
    [5,7,1,2,9],
    [1,3,5,6,6]
])
 
# 返回x中0元素的索引
index0 = torch.where(x==0) # 輸入?yún)?shù)為1個
 
print(index0,'\n', type(index0))

結果

>>> index0
(tensor([0, 1, 2]), tensor([3, 4, 2])) 
 
>>> type(index0)
<class 'tuple'>

其中[0, 1, 2]是0元素坐標的行索引，[3, 4, 2]是0元素坐標的列索引，注意，最終得到的是tuple類型的返回值，元組中包含了tensor

2.np.where()

np.where()用法與torch.where()用法類似，也包括兩種用法，但是不同的是輸入值類型和返回值的類型

代碼示例

np.where(condition, x, y)和np.where(condition)，輸入x,y可以為非tensor

代碼

import torch
import numpy as np
x = torch.tensor([
    [1,2,3,0,6],
    [4,6,2,1,0],
    [4,3,0,1,1]
])
y = torch.tensor([
    [0,5,1,4,2],
    [5,7,1,2,9],
    [1,3,5,6,6]
])
 
arr1 = np.where(x>=3, x, y) # 輸入?yún)?shù)為3個
 
index0 = torch.where(x==0) # 輸入?yún)?shù)為1個
 
print(arr1,'\n',type(arr1))
print(index1,'\n', type(index1))
 

結果

>>> arr1
[[0 5 3 4 6]
 [4 6 1 2 9]
 [4 3 5 6 6]]
 
>>> type(arr1)
<class 'numpy.ndarray'>
 
>>> index1
(array([0, 1, 2]), array([3, 4, 2])) 
 
>>> type(index1)
<class 'tuple'>

注意，np.where()和torch.where()的返回值類型不同

3.np.argwhere(condition)

尋找符合contion的元素索引

代碼示例

代碼

import torch
import numpy as np
x = torch.tensor([
    [1,2,3,0,6],
    [4,6,2,1,0],
    [4,3,0,1,1]
])
y = torch.tensor([
    [0,5,1,4,2],
    [5,7,1,2,9],
    [1,3,5,6,6]
])
 
 
index2 = np.argwhere(x==0) # 尋找元素為0的索引
 
print(index2,'\n', type(index2))

結果

>>> index2
tensor([[0, 1, 2],
        [3, 4, 2]]) 
 
>>> type(index2)
<class 'torch.Tensor'>

注意返回值的類型

總結

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論