Как трехмерный тензор индексируется двумя двумерными тензорами? - PullRequest
Новая
       14

Как трехмерный тензор индексируется двумя двумерными тензорами?

0 голосов
/ 15 сентября 2018

Вот снимок со строки 15-20 в DIM 

def random_permute(X):
    X = X.transpose(1, 2)
    b = torch.rand((X.size(0), X.size(1))).cuda()
    idx = b.sort(0)[1]
    adx = torch.range(0, X.size(1) - 1).long()
    X = X[idx, adx[None, :]].transpose(1, 2)

    return X

, где X - тензор размера [64, 64, 128], idx тензор размера [64, 64], adx тензор размера [64]. Как работает X = X[idx, adx[None, :]]? Как мы можем использовать два 2d-тензора для индексации 3d-тензора? Что на самом деле происходит с X после индексации?

Ответы [ 2 ]

0 голосов
/ 22 сентября 2018

Вещи станут более понятными, если мы рассмотрим конкретный пример меньшего размера.Пусть 

x = np.arange(8).reshape(2, 2, 2)
b = np.random.rand(2, 2)
idx = b.argsort(0) # e.g. idx=[[1, 1], [0, 0]]
adx = np.arange(2)[None, :] # [[0, 1]]
y = x[idx, adx] # implicitly expanding 'adx' to [[0, 1], [0, 1]]

В этом примере у нас будет y как 

y[0, 0] = x[idx[0, 0], adx[0, 0]]=x[1, 0]
y[0, 1] = x[idx[0, 1], adx[0, 1]]=x[1, 1]
y[1, 0] = x[idx[1, 0], adx[1, 0]]=x[0, 0]
...

. Может быть полезно увидеть, как мы делаем то же самое в тензорном потоке: 

d0, d1, d2 = x.shape.as_list()
b = np.random.rand(d0, d1)
idx = np.argsort(b, 0)
idx = idx.reshape(-1)
adx = np.arange(0, d1)
adx = np.tile(adx, d0)

y = tf.reshape(tf.gather_nd(x, zip(idx, adx)), (d0, d1, d2))
0 голосов
/ 17 сентября 2018

По моему предположению X должен быть 3D-тензором, поскольку обычно он представляет собой серию обучающих данных.

Что касается функциональности этой функции, она случайным образом переставляет тензор входных данных X и делает это, используя следующие шаги:

  • Сначала он инициализирует тензор b значениями, выбранными из равномерного распределения.
  • Далее этот тензор сортируется по измерению 0, а индексы сортировки выводятся в тензор idx.
  • Тензор adx - это просто целочисленный тензор значений в диапазоне от 0 до 63.

Теперь в строке ниже происходит вся магия: 

X[idx, adx[None, :]].transpose(1, 2)

Мы используем индексы, которые мы получили до idx и adx (adx[None, :] - просто двумерный вектор строки). Получив это, мы транспонируем оси 1 и 2 точно так же, как мы делали в начале функции в строке: 

X = X.transpose(1, 2)

Вот надуманный пример, для лучшего понимания: 

# our input tensor
In [51]: X = torch.rand(64, 64, 32)

In [52]: X = X.transpose(1, 2)

In [53]: X.shape
Out[53]: torch.Size([64, 32, 64])

In [54]: b = torch.rand((X.size(0), X.size(1)))

# sort `b` which returns a tuple and take only indices
In [55]: idx = b.sort(0)[1]

In [56]: idx.shape
Out[56]: torch.Size([64, 32])

In [57]: adx = torch.arange(0, X.size(1)).long()

In [58]: adx.shape
Out[58]: torch.Size([32])

In [59]: X[idx, adx[None, :]].transpose(1, 2).shape
Out[59]: torch.Size([64, 64, 32])

Важно отметить, что на последнем шаге мы получили ту же форму, что и форму входного тензора, равного (64, 64, 32).

...