Use a generator to perform operation on a matrix in Python
I have a similarity matrix (torch tensor) which is a cosine similarity matrix between two matrix (source and target). From the matrix I need to obtain the sum of the distance between the top nearest neighbor of each source and target.  

Then fillup two defaultdicts using the computed values above as shown in the code snippet below

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() )



Is there a way I can optimize the above code, either using a generator without having to create nearestSrc, sumDistSource, nearestTgt, sumDistTarget explicitly thus requiring less memory?  Or can I also reduce the double loop?

I don't think it's necessary to save memory here. Let's say the shape of matx is [n x m], then the nearestSrc/Tgt and sumDistSource/Target tensors will contain no more than 2 * (n + m), memory consumption of which is almost ignorable compared to matx. Besides, I don't think PyTorch provides an API to generate top-k elements on-the-fly, and it would be hard to implement a differentiable or GPU-optimized version of that.



The double-loop can be optimized, although I'm a bit confused by what you're doing here. It seems that the values you're computing for src2tgt[src[i]][tgt[j]] and tgt2src[tgt[j]][src[i]] are exactly the same. Also, I don't think storing these in a nested dictionary is a good idea, for two reasons:


src[i] and tgt[j] are floating-point tensors, and interestingly, the hash value for PyTorch tensors does not depend on tensor values. In fact, it is equivalent to the id function. See more discussions here. This means that two tensors with the same value would still be stored as different keys in the dictionary. Besides, it's also probably not a good idea to use floating-point numbers as keys, because testing equality for floating-point numbers often require special care. For more information on this topic, consider reading this very helpful blog.
Taking values out of a tensor and storing them in another structure would prevent further optimizations. Operations on tensors can benefit from the highly-optimized PyTorch functions, which scales sublinearly on GPUs (i.e., a 2x increase in data size leads to a &lt;2x increase in compute time). It's often more desirable to do stuff in tensor form as much as you can.


Thus, we can optimize your code by first computing all the values you need to store in src2tgt by a batch tensor operation, and then storing them into the dictionary.

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()


However, the speedup would be small as long as you still use a dictionary. I would encourage you to consider other ways to store the data.
