如何将张量插入到 pytorch 中的另一个张量中

debugcn 发表于 Dev

Tong

我有 pytorch Tensor 的 shape (batch_size, step, vec_size)，例如， a Tensor(32, 64, 128)，我们称之为 A 。

我有另一个Tensor(batch_size, vec_size)，例如Tensor(32, 128)，我们称其为 B。

我想将 B 插入到 A 轴 1 的某个位置。插入位置在Tensor(batch_size)名为 P 的a 中给出。

我知道 pytorch 中没有空张量（如空列表），因此，我将 A 初始化为零，并在 A 轴 1 的某个位置添加 B。

A = Variable(torch.zeros(batch_size, step, vec_size))

我正在做的是：

for i in range(batch_size):
    pos = P[i]
    A[i][pos] = A[i][pos] + B[i]

但我收到一个错误：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

然后，我在循环中对每个 A 进行克隆：

for i in range(batch_size):
    A_clone =  A.clone()
    pos = P[i]
    A_clone[i][pos] = A_clone[i][pos] + B[i]

这对于 autograd 来说非常慢，我想知道是否有更好的解决方案？谢谢你。

乌曼古普塔

您可以使用掩码代替克隆。

看下面的代码

# setup
batch, step, vec_size = 64, 10, 128 
A = torch.rand((batch, step, vec_size))
B = torch.rand((batch, vec_size))
pos = torch.randint(10, (64,)).long()

# computations
# create a mask where pos is 0 if it is to be replaced
mask = torch.ones( (batch, step)).view(batch,step,1).float()
mask[torch.arange(batch), pos]=0

# expand B to have same dimension as A and compute the result
result = A*mask + B.unsqueeze(dim=1).expand([-1, step, -1])*(1-mask)

这样您也可以避免使用 for 循环和克隆。

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。