Используйте генератор для выполнения операции над матрицей в Python

Question

У меня есть матрица подобия (тензор факела), которая представляет собой косинусную матрицу подобия между двумя матрицами (источником и целью). Из матрицы мне нужно получить сумму расстояния между верхним ближайшим соседом каждого источника и цели.

Затем заполните два defaultdicts, используя приведенные выше вычисленные значения, как показано в фрагменте кода ниже

import torch
from collections import defaultdict
src2tgt = defaultdict(dict)
tgt2src = defaultdict(dict)

#similarity matrix between source and target matrix
matx = torch.Tensor([[3,2,1,7],[1,1,0,8],[0,7,1,0],[2,0,0,0],[1,5,2,1]])

#the src and tgt 
src = torch.LongTensor([[1,1],[1,2],[1,3],[1,4],[1,5]])
tgt = torch.LongTensor([[2,1],[2,2],[2,3],[2,4]])
#the data above are dummy, in my actual code, they are generated by a process
similarities=[]

#I need a kinda of the sum of nearest neighbor
k = 2
nearestSrc = torch.topk(matx, k, dim=1, largest=True, sorted=False, out=None)
sumDistSource = torch.sum(nearestSrc[0], 1)
nearestTgt = torch.topk(matx, k, dim=0, largest=True, sorted=False, out=None)
sumDistTarget = torch.sum(nearestTgt[0], 0)


#finally fill default dictionary of source2target and target2source
for i in range(matx.shape[0]):
  for j in range(matx.shape[1]):
    src2tgt[src[i]][tgt[j]] = matx[i][j].tolist() / (sumDistSource[i].tolist() + sumDistTarget[j].tolist())
    tgt2src[tgt[j]][src[i]] = matx[i][j].tolist() / (sumDistTarget[j].tolist() + sumDistSource[i].tolist())
    similarities.append(matx[i][j].tolist() )

Есть ли способ оптимизировать приведенный выше код, используя генератор без создать nearestSrc, sumDistSource, nearestTgt, sumDistTarget, явно требуя меньше памяти? Или я могу также уменьшить двойной l oop?

Answer 1

Я не думаю, что здесь необходимо экономить память. Скажем, форма matx равна [n x m], тогда тензоры nearestSrc/Tgt и sumDistSource/Target будут содержать не более 2 * (n + m), потребление памяти в которых почти игнорируется по сравнению с matx. Кроме того, я не думаю, что PyTorch предоставляет API для генерации элементов top-k на лету, и было бы сложно реализовать его дифференцируемую или оптимизированную для GPU версию.

---

double-l oop можно оптимизировать, хотя меня немного смущает то, что вы здесь делаете. Кажется, что значения, которые вы вычисляете для src2tgt[src[i]][tgt[j]] и tgt2src[tgt[j]][src[i]], абсолютно одинаковы. Кроме того, я не думаю, что хранить их во вложенном словаре - хорошая идея по двум причинам:

1. src[i] и tgt[j] являются тензорами с плавающей точкой, и что интересно, значение ha sh для тензоров PyTorch не зависит от значений тензоров . Фактически это эквивалентно функции id. Больше обсуждений здесь . Это означает, что два тензора с одинаковым значением будут по-прежнему храниться в словаре в качестве разных ключей. Кроме того, вероятно, также не очень хорошая идея использовать числа с плавающей точкой в ​​качестве ключей, потому что проверка равенства для чисел с плавающей точкой часто требует особого внимания. Для получения дополнительной информации по этой теме c рассмотрите этот очень полезный блог .

2. Извлечение значений из тензора и сохранение их в другой структуре предотвратит дальнейшее оптимизаций. Операции над тензорами могут получить выгоду от высокооптимизированных функций PyTorch, которые масштабируются сублинейно на графических процессорах (то есть увеличение размера данных в 2 раза приводит к увеличению времени вычислений в <2 раза). Часто более желательно делать вещи в тензорной форме как можно больше. </p>

Таким образом, мы можем оптимизировать ваш код, сначала вычислив все значения, которые необходимо сохранить в src2tgt, с помощью пакетная тензорная операция с последующим сохранением их в словаре.

norm = (sumDistSource.unsqueeze(-1).expand(-1, matx.size(1)) +
        sumDistTarget.unsqueeze(0).expand(matx.size(0), -1))
s2t = (matx / norm).tolist()

src_vals = src.tolist()
tgt_vals = tgt.tolist()
for i, s in enumerate(src_vals):
    s = tuple(s)  # convert to tuples because lists are not hashable
    for j, t in enumerate(tgt_vals):
        t = tuple(t)
        src2tgt[s][t] = s2t[i][j]
        tgt2src[t][s] = s2t[i][j]
similarities = matx.flatten().tolist()

Однако ускорение будет небольшим, если вы все еще используете словарь. Я бы посоветовал вам рассмотреть другие способы хранения данных.