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.inputwill 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]])