Логика выбора тензорного входа c Дивергенция в PyTorch & Numpy

Question

Описание

Я настраиваю torch.Tensor для маскировки. При попытке выбора записей по индексам оказывается, что поведение между использованием numpy.ndarray и torch.Tensor для хранения данных индекса отличается. Я хотел бы получить доступ к дизайну как в рамках, так и в связанных документах, объясняющих разницу.

Шаги для репликации

Среда

Pytorch 1.3 в контейнере из официального выпуска: pytorch/pytorch:1.3-cuda10.1-cudnn7-devel

Пример

Скажем, мне нужно установить mask как torch.Tensor объект с формой [3,3,3] и установить значения в записях (0,0,1) & (1,2,0) в 1. Код ниже объясняет разницу.

mask = torch.zeros([3,3,3])
indices = torch.tensor([[0, 1],
                        [0, 2],
                        [1, 0]])

mask[indices.numpy()] = 1 # Works
# mask[indices] = 1 # Incorrect result

Я заметил, что при использовании mask[indices.numpy()] новой torch.Tensor формы [2], тогда как mask[indices] возвращает новую torch.Tensor формы [3, 2, 3, 3] , что говорит о разнице в логике тензорной нарезки c.

Answer 1

Вы получаете разные результаты, потому что именно так индексация реализована в Pytorch. Если вы передадите массив в качестве индекса, он будет «распакован». Например:

indices = torch.tensor([[0, 1], [0, 2], [1, 0]])

mask = torch.arange(1,28).reshape(3,3,3)

# tensor([[[ 1,  2,  3],
#          [ 4,  5,  6],
#          [ 7,  8,  9]],

#         [[10, 11, 12],
#          [13, 14, 15],
#          [16, 17, 18]],

#         [[19, 20, 21],
#          [22, 23, 24],
#          [25, 26, 27]]])

mask[indices.numpy()] эквивалентно mask[[0, 1], [0, 2], [1, 0]], т.е. элементы i-й строки indices.numpy() используются для выбора элементов mask вдоль i-й оси , Таким образом, он возвращает tensor([mask[0,0,1], mask[1,2,0]]), то есть tensor([2, 16]).

С другой стороны, при передаче тензора в качестве индекса (я не знаю точную причину такого различия между массивами и тензорами для индексации), он не "распаковывается" как массив, и элементы i-й строки тензора indices используются для выбора элементов mask вдоль оси-0. То есть mask[indices] эквивалентен mask[[[0, 1], [0, 2], [1, 0]], :, :]

>>> mask[ind]

tensor([[[[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16, 17, 18]]],


        [[[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26, 27]]],


        [[[10, 11, 12],
          [13, 14, 15],
          [16, 17, 18]],

         [[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]]]])

, который в основном tensor(mask[[0,1], :, :], mask[[0,2],: ,:], mask[[1,0], :, :]) и имеет форму indices.shape + mask[0,:,:].shape == (3,2,3,3). Таким образом, целые «листы» отбираются и укладываются в новые измерения. Обратите внимание, что это не новый тензор, а особое представление mask. Следовательно, если вы присваиваете mask[indices] = 1 с этим конкретным indices, то все элементы mask станут равны 1.