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

Question

我有一个形状为(batch_size, step, vec_size)的pytorch Tensor，例如Tensor(32, 64, 128)，我们称之为A。

我还有另一个Tensor(batch_size, vec_size)，例如Tensor(32, 128)，我们就叫它B吧。

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

我知道 pytorch 中没有 Empty tensor（如空列表），因此，我将 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 each：

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

这对于autograd来说很慢，不知道有没有更好的解决方案？谢谢。

Answer 1

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

请看下面的代码

# 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 循环和克隆。