Как я могу вставить Тензор в другой Тензор в pytorch

Question

У меня есть тензор Pytorch с формой (batch_size, step, vec_size), например, Tensor(32, 64, 128), назовем его A.

У меня есть другой Tensor(batch_size, vec_size), например, Tensor(32, 128), назовем его B.

Я хочу вставить B в определенную позицию на оси 1 из A. Позиции вставки даны в Tensor(batch_size) с именем P.

Я понимаю, что пустого тензора не существует (например,пустой список) в pytorch, поэтому я инициализирую A как нули и добавляю B в определенную позицию на оси 1 в A.

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]

Это очень медленноИнтересно, есть ли лучшие решения?Спасибо.

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)

Таким образом, вы также избегаете использования для циклов и клонирования.