torch.Tensor.index_add_¶
- Tensor.index_add_(dim, index, source, *, alpha=1) 张量¶
将
alpha乘以source的元素累加到self张量中,按照index中给定的顺序添加到索引中。例如,如果dim == 0,index[i] == j,并且alpha=-1,那么source的第i行将从self的第j行中减去。dim维度中的source必须与index的长度(必须是一个向量)相同,并且所有其他维度必须与self匹配,否则将引发错误。对于一个3维张量,输出如下所示:
self[index[i], :, :] += alpha * src[i, :, :] # 如果 dim == 0 self[:, index[i], :] += alpha * src[:, i, :] # 如果 dim == 1 self[:, :, index[i]] += alpha * src[:, :, i] # 如果 dim == 2
注意
当在CUDA设备上使用张量时,此操作可能会表现出不确定性行为。更多信息请参见可重复性。
- Parameters
- Keyword Arguments
alpha (数字) –
source的标量乘数
示例:
>>> x = torch.ones(5, 3) >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) >>> index = torch.tensor([0, 4, 2]) >>> x.index_add_(0, index, t) tensor([[ 2., 3., 4.], [ 1., 1., 1.], [ 8., 9., 10.], [ 1., 1., 1.], [ 5., 6., 7.]]) >>> x.index_add_(0, index, t, alpha=-1) tensor([[ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.], [ 1., 1., 1.]])