Предварительное индексирование в Pytorch, чтобы избавиться от вложенных циклов for - PullRequest
Новая
       57

Предварительное индексирование в Pytorch, чтобы избавиться от вложенных циклов for

1 голос
/ 13 февраля 2020

У меня есть ситуация, для которой я использую вложенные циклы for, но я хочу знать, есть ли более быстрый способ сделать это с помощью некоторого расширенного индексирования в Pytorch.

У меня есть тензор с именем t: 

t = torch.randn(3,8)
print(t)
tensor([[-1.1258, -1.1524, -0.2506, -0.4339,  0.8487,  0.6920, -0.3160, -2.1152],
        [ 0.4681, -0.1577,  1.4437,  0.2660,  0.1665,  0.8744, -0.1435, -0.1116],
        [ 0.9318,  1.2590,  2.0050,  0.0537,  0.6181, -0.4128, -0.8411, -2.3160]])

Я хочу создать новый тензор, который индексирует значения из t. Допустим, эти индексы хранятся в переменной indexes

indexes = [[(0, 1, 4, 5), (0, 1, 6, 7), (4, 5, 6, 7)],
           [(2, 3, 4, 5)],
           [(4, 5, 6, 7), (2, 3, 6, 7)]]

Каждый внутренний кортеж в indexes представляет четыре индекса, которые должны быть взяты из строки.

В качестве примера на основе по этим индексам мой вывод будет тензор измерения 6x4 (6 - общее количество кортежей в indexes, а 4 соответствует одному значению в кортеже)

Например, это то, что я хочу сделать : 

#counting the number of tuples in indexes
count_instances = sum([1 for lst in indexes for tupl in lst])

#creating a zero output matrix 
final_tensor = torch.zeros(count_instances,4)

final_tensor[0] = t[0,indexes[0][0]]
final_tensor[1] = t[0,indexes[0][1]]
final_tensor[2] = t[0,indexes[0][2]]
final_tensor[3] = t[1,indexes[1][0]]
final_tensor[4] = t[2,indexes[2][0]]
final_tensor[5] = t[2,indexes[2][1]]

Окончательный вывод выглядит следующим образом: print (final_tensor) 

tensor([[-1.1258, -1.1524,  0.8487,  0.6920],
        [-1.1258, -1.1524, -0.3160, -2.1152],
        [ 0.8487,  0.6920, -0.3160, -2.1152],
        [ 1.4437,  0.2660,  0.1665,  0.8744],
        [ 0.6181, -0.4128, -0.8411, -2.3160],
        [ 2.0050,  0.0537, -0.8411, -2.3160]])

Я создал функцию build_tensor (показано ниже) для достижения этого с помощью вложенных циклов for, но я хочу знать, есть ли более быстрый способ сделать это с помощью простой индексации в Pytorch. Я хочу более быстрый способ сделать это, потому что я делаю эту операцию в сотни раз с большими индексами и размерами t.

Любая помощь? 

def build_tensor(indexes, t):
    #count tuples
    count_instances = sum([1 for lst in indexes for tupl in lst])
    #create a zero tensor
    final_tensor = torch.zeros(count_instances,4)
    final_tensor_idx = 0

    for curr_idx, lst in enumerate(indexes):
        for tupl in lst:
            final_tensor[final_tensor_idx] = t[curr_idx,tupl]
            final_tensor_idx+=1
    return final_tensor

1 Ответ

1 голос
/ 13 февраля 2020

Вы можете упорядочить индексы в двумерные массивы, а затем выполнить индексацию за один кадр следующим образом: 

rows = [(row,)*len(index_tuple) for row, row_indices in enumerate(indexes) for index_tuple in row_indices]
columns = [index_tuple for row_indices in indexes for index_tuple in row_indices]
final_tensor = t[rows, columns]
...