Pytorch torch.Where 和 torch.Gather

torch.Where

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

返回一个tensor，tensor 的元素从x或者y中的元素选择，选择依据是条件矩阵condition。

要求：condition、x，y可以broadcastable到相同的shape。返回值也是相同shape的矩阵

condition：

参数：

• condition (BoolTensor) – 条件矩阵，当元素为 True 时，从x中选择对应的元素，否则填入 y 中对应的元素。

• x (Tensor or Scalar) – 标量值 或者张量

• y (Tensor or Scalar) – 标量值 或者张量

返回值

• 与condition, x, y broadcast以后相同shape的张量

返回类型

• Tensor

举例说明：

>>> 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.gather

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

沿由dim 指定的轴收集轴上面的由index索引得到的元素的值。

例如对于一个3D的tensor：

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

输入和索引必须具有相同的维数。

参数：

• input (Tensor) – 源 tensor

• dim (int) – 收集的元素的所在的轴

• index (LongTensor) – 要收集的元素的索引

关键字参数：

• sparse_grad (bool, optional) – If True, gradient w.r.t. input will be a sparse tensor.

• out (Tensor, optional) – 输出tensor

举例说明：

>>> t = torch.tensor([[1, 2], [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
tensor([[ 1,  1],
        [ 4,  3]])