PyTorch: применить отображение по одноэлементному измерению тензора - PullRequest
Новая
       29

PyTorch: применить отображение по одноэлементному измерению тензора

1 голос
/ 21 апреля 2019

Боюсь, название не очень описательное, но я не мог придумать лучшего. По сути, моя проблема заключается в следующем:

У меня есть тензор pytorch формы (n, 1, h, w) для произвольных целых чисел n, h и w (в моем конкретном случае этот массив представляет серию изображений в градациях серого размера h x w).

У меня также есть другой тензор формы (m, 2), который отображает каждое возможное значение в первом массиве (то есть первый массив может содержать значения от 0 до m - 1) в некоторый набор значений. Я хотел бы применить это отображение к первому массиву, чтобы получить массив формы (n, 2, h, w).

Надеюсь, это несколько понятно, мне трудно это выразить словами, вот пример кода (но обратите внимание, что это не является сверхинтуитивным из-за использования четырехмерных массивов): 

import torch

m = 18

# could also be arbitrary tensor with this shape with values between 0 and m - 1
a = torch.arange(m).reshape(2, 1, 3, 3)

# could also be arbitrary tensor with this shape
b = torch.LongTensor(
    [[11, 17, 9, 6, 5, 4, 2, 10, 3, 13, 14, 12, 7, 1, 15, 16, 8, 0],
     [11, 8, 4, 14, 13, 12, 16, 1, 5, 17, 0, 10, 7, 15, 9, 6, 2, 3]]).t()

# I probably have to do this and the permute/reshape, but how?
c = b.index_select(0, a.flatten())

# ...

# another approach that I think works (but I'm not really sure why, I found this
# more or less by trial and error). I would ideally like to find a 'nicer' way
# of doing this
c = torch.stack([
    b.index_select(0, a_.flatten()).reshape(3, 3, 2).permute(2, 0, 1)
    for a_ in a
])

# the end result should be:
#[[[[11, 17,  9],
#   [ 6,  5,  4],
#   [ 2, 10,  3]],
#
#  [[11,  8,  4],
#   [14, 13, 12],
#   [16,  1,  5]]],
#
#
# [[[13, 14, 12],
#   [ 7,  1, 15],
#   [16,  8,  0]],
#
#  [[17,  0, 10],
#   [ 7, 15,  9],
#   [ 6,  2,  3]]]]

Как я могу выполнить это преобразование эффективным образом? (В идеале не использовать дополнительную память). В NumPy это легко может быть достигнуто с помощью np.apply_along_axis, но, кажется, не существует Pytorch эквивалентный этому.

1 Ответ

1 голос
/ 21 апреля 2019

Вот один из способов использования нарезки, укладки и изменения формы на основе вида: 

In [239]: half_way = b.shape[0]//2

In [240]: upper_half = torch.stack((b[:half_way, :][:, 0], b[:half_way, :][:, 1]), dim=0).view(-1, 3, 3)
In [241]: lower_half = torch.stack((b[half_way:, :][:, 0], b[half_way:, :][:, 1]), dim=0).view(-1, 3, 3)

In [242]: torch.stack((upper_half, lower_half))
Out[242]: 
tensor([[[[11, 17,  9],
          [ 6,  5,  4],
          [ 2, 10,  3]],

         [[11,  8,  4],
          [14, 13, 12],
          [16,  1,  5]]],


        [[[13, 14, 12],
          [ 7,  1, 15],
          [16,  8,  0]],

         [[17,  0, 10],
          [ 7, 15,  9],
          [ 6,  2,  3]]]])

Некоторые предостережения в том, что это будет работать только для n=2. Однако это в 1,7 раза быстрее, чем ваш подход, основанный на циклах, но требует больше кода.

Вот более обобщенный подход , который масштабируется до любого положительного целого числа n: 

In [327]: %%timeit
     ...: block_size = b.shape[0]//a.shape[0]
     ...: seq_of_tensors = [b[block_size*idx:block_size*(idx+1), :].permute(1, 0).flatten().reshape(2, 3, 3).unsqueeze(0)  for idx in range(a.shape[0])]
     ...: torch.cat(seq_of_tensors)
     ...: 
23.5 µs ± 460 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Вы также можете использовать view вместо изменения формы: 

block_size = b.shape[0]//a.shape[0]
seq_of_tensors = [b[block_size*idx:block_size*(idx+1), :].permute(1, 0).flatten().view(2, 3, 3).unsqueeze(0)  for idx in range(a.shape[0])]
torch.cat(seq_of_tensors)
# outputs
tensor([[[[11, 17,  9],
          [ 6,  5,  4],
          [ 2, 10,  3]],

         [[11,  8,  4],
          [14, 13, 12],
          [16,  1,  5]]],


        [[[13, 14, 12],
          [ 7,  1, 15],
          [16,  8,  0]],

         [[17,  0, 10],
          [ 7, 15,  9],
          [ 6,  2,  3]]]])

Примечание : обратите внимание, что я все еще использую понимание списка, поскольку мы должны равномерно разделить наш тензор b, чтобы переставить, сгладить, изменить форму, отжать, а затем объединить / сложить вдоль измерения 0. Это все еще немного быстрее, чем мое решение выше.

...