How to efficiently calculate distance matrix in pytorch for two sets 3D tensors with different sizes?
I have tensors X of shape BxCxHxW and Y of shape NxCxHxW. B is batch size, C is channels, H is height, W is width, and N will be constant for any batch. Basically I want the BxN distance matrix of distances between a set of B images and another set of N images.
I tried using torch.cdist by reshaping X as 1xBx(C*H*W) and Y as 1xNx(C*H*W) by unsqueezing a dimension and flattening the last 3 channels, but I did a sanity check and got wrong answers with this method.
I want L2 distance.

According to the documentation page for torch.cdist, the two inputs and outputs are shaped in the following manner: x1: (B, P, M), x2: (B, R, M), and output: (B, P, R).
To match your case: B=1, P=B, R=N, while M=C*H*W (i.e. flattened). As you just explained.
So you are basically going for:
&gt;&gt;&gt; torch.cdist(X[None].flatten(2), Y[None].flatten(2))


If you're not convinced, you can check with the following method:
&gt;&gt;&gt; dist = []
&gt;&gt;&gt; for x in X:
...    for y in Y:
...       dist.append((x-y).norm())

And compare the torch.cdist result with torch.tensor(dist).reshape(len(X), len(Y)).
