4D tensor correlation


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

Hi Rane!

There most likely is.

torch.corrcoef() does not seem to operate on batches of matrices, so
the approach would be to write your own batched version of corrcoef().

In your previous thread, we looked at how to write a batched version of
your cov() function. I would expect that you would be able to do something
broadly similar with corrcoef().

Give it a try, and if you have questions, please post a fully-self-contained,
runnable script, together with its output, that uses for-loops to perform your
desired computation followed by your attempt at a loop-free version.


K. Frank