Hi,
I’m having trouble writing code for the following scenario:
I have two 4D tensors <X, Y>
of shape (b, c, h, w)
. For every c
I would like to compute the cross-correlation matrix between X
and Y
based on the i
-th index of w
(i.e. correlation based on the i
-th columns of each (h,w)
matrix).
So for example, if I choose i = 1
, the following code calculates the correlation coefficient between column 1
of the first two matrixes of c
, from the first element of the batch of X
and Y
(i.e. index 0
of c
and b
):
X = torch.randn((3, 2, 4, 5))
Y = torch.randn((3, 2, 4, 5))
i = 1
X_col_i = = X[0, 0, :, i]
Y_col_i = Y[0, 0, :, i]
stacked_cols = torch.vstack((col_i_X, Y_col_i))
coeff = torch.corrcoef(stacked_cols)[0,1]
Is there any way to vectorize these operations while avoiding a nested for loop for b
and c
?
The output shape should be (b, c, c)
: a (c, c)
cross-correlation matrix for each element in the batch.
Thank you