Python PyTorch where用法及代碼示例

本文簡要介紹python語言中 torch.where 的用法。

用法:

torch.where(condition, x, y) → Tensor

參數：

• condition(BoolTensor) -當 True(非零)時，產生 x，否則產生 y

• x(Tensor或者標量) -值(如果 x 是標量)或在 condition 為 True 的索引處選擇的值

• y(Tensor或者標量) -值(如果 y 是標量)或在 condition 為 False 的索引處選擇的值

返回：

形狀等於 condition 、 x 、 y 的廣播形狀的張量

返回類型：

Tensor

返回從 x 或 y 中選擇的元素的張量，具體取決於 condition 。

操作定義為：

\text{out}_i = \begin{cases} \text{x}_i & \text{if } \text{condition}_i \\ \text{y}_i & \text{otherwise} \\ \end{cases}

注意

張量 condition 、 x 、 y 必須是可廣播的。

注意

當前有效的標量和張量組合是 1. 浮點數 dtype 和 torch.double 的標量 2. 積分 dtype 和 torch.long 的標量 3. 複雜 dtype 和 torch.complex128 的標量

例子：

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
        [ 0.3898, -0.7197],
        [ 0.0478, -0.1657]])
>>> torch.where(x > 0, x, y)
tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])
>>> x = torch.randn(2, 2, dtype=torch.double)
>>> x
tensor([[ 1.0779,  0.0383],
        [-0.8785, -1.1089]], dtype=torch.float64)
>>> torch.where(x > 0, x, 0.)
tensor([[1.0779, 0.0383],
        [0.0000, 0.0000]], dtype=torch.float64)

torch.where(condition) → tuple of LongTensor

torch.where(condition) 與 torch.nonzero(condition, as_tuple=True) 相同。

注意

另見 torch.nonzero() 。