Pytorch select значения из последнего тензорного измерения с индексами другого тенора с меньшим измерением

Question

У меня есть тензор a с тремя измерениями. Первое измерение соответствует размеру мини-пакета, второе - длине последовательности, а третье - измерению объекта. Например.,

>>> a = torch.arange(1, 13, dtype=torch.float).view(2,2,3)  # Consider the values of a to be random
>>> a
tensor([[[ 1.,  2.,  3.],
         [ 4.,  5.,  6.]],

        [[ 7.,  8.,  9.],
         [10., 11., 12.]]])

У меня есть второй, двумерный тензор. Его первое измерение соответствует размеру мини-пакета, а второе - длине последовательности. Содержит значения в диапазоне индексов третьего измерения a. Третье измерение a s имеет размер 3, поэтому b может содержать значения 0, 1 или 2. Например,

>>> b = torch.LongTensor([[0, 2],[1,0]])
>>> b
tensor([[0, 2],
        [1, 0]])

Я хочу получить тензор c, который имеет форму b и содержит все значения a, на которые ссылается b. В верхнем сценарии я хотел бы иметь:

c = torch.empty(2,2)
c[0,0] = a[0, 0, b[0,0]]
c[1,0] = a[1, 0, b[1,0]]
c[0,1] = a[0, 1, b[0,1]]
c[1,1] = a[1, 1, b[1,1]]

>>> c
tensor([[ 1.,  5.],
        [ 8., 10.]])

Как быстро создать тензор c? Кроме того, я также хочу, чтобы c был дифференцируем (мог использовать .backprob()) Я не слишком знаком с pytorch, поэтому я не уверен, существует ли дифференцируемая версия этого.

В качестве альтернативы вместо c, имеющей ту же форму, что и b, я мог бы также использовать c с той же формой a, имеющей только нули, но в местах, на которые ссылается b из них. Тогда я мог бы умножить a и c, чтобы получить дифференцируемый тензор.

Как следует:

c = torch.zeros(2,2,3, dtype=torch.float)
c[0,0,b[0,0]] = 1
c[1,0,b[1,0]] = 1
c[0,1,b[0,1]] = 1
c[1,1,b[1,1]] = 1

>>> a*c
tensor([[[ 1.,  0.,  0.],
         [ 0.,  5.,  0.]],

        [[ 0.,  8.,  0.],
         [10.,  0.,  0.]]])

Answer 1

Позволяет сначала объявить необходимые переменные: (обратите внимание requires_grad при инициализации a, мы будем использовать его для обеспечения дифференцируемости)

a = torch.arange(1,13,dtype=torch.float32,requires_grad=True).reshape(2,2,3)
b = torch.LongTensor([[0, 2],[1,0]])

Позволяет изменить форму a и размеры мини-пакета и последовательности сквоша:

temp = a.reshape(-1,3)

, поэтому temp теперь выглядит следующим образом:

tensor([[ 1.,  2.,  3.],
    [ 4.,  5.,  6.],
    [ 7.,  8.,  9.],
    [10., 11., 12.]], grad_fn=<AsStridedBackward>)

Обратите внимание, что теперь каждое значение b может использоваться в каждой строке temp для получения желаемого результата.Теперь мы делаем:

c = temp[range(len(temp )),b.view(-1)].view(b.size())

Обратите внимание, как мы индексируем temp, range(len(temp )) для выбора каждой строки и 1D b, то есть b.view(-1) для получения соответствующих столбцов.Наконец, .view(b.size()) приводит этот массив к тому же размеру, что и b.

Если мы печатаем c сейчас:

tensor([[ 1.,  6.],
    [ 8., 10.]], grad_fn=<ViewBackward>)

Наличие grad_fn=.. показывает, что cтребует градиента, т.е. его дифференцируемо.