Custom loss - correlation by column

username: Rane90
cooked: <p>Hi,</p>
<p>I’ve been having trouble computing my desired loss: I have two input tensors &lt;X, Y&gt;, with shapes: (batch_size, num_rows, num_columns). The loss I want to compute is the correlation between the respective &lt;X, Y&gt; columns for each one of the samples.</p>
<p>So far I’ve been iterating through each sample:</p>
<pre><code class="lang-auto">def cov(x, y):
    x_bar = x - x.mean(axis=0)
    y_bar = y - y.mean(axis=0)
    N = x_bar.shape[0]
    return (y_bar.T @ x_bar).T / (N - 1)


class Corr(nn.Module):
    def __init__(self, device, eps=1e-5):
        super(Corr, self).__init__()
        self.device = device
        self.eps = eps

    def forward(self, x, y, params):

        loss = 0
        for i in range(x.shape[0]):
            xi = x[i, :, :]
            yi = y[i, :, :]
            C_yy = cov(yi, yi)
            C_yx = cov(yi, xi)
            C_xx = cov(xi, xi)

            C_yy = C_yy + \
                torch.eye(C_yy.shape[0], device=self.device) * self.eps
            C_xx = C_xx + \
                torch.eye(C_xx.shape[0], device=self.device) * self.eps

            M = torch.linalg.multi_dot([torch.inverse(C_yy),
                                        C_yx,
                                        torch.inverse(C_xx),
                                        C_yx.T])

            loss += torch.trace(M)

        return (loss/x.shape[0])
</code></pre>
<p>Is there a faster way to implement this?</p>

username: KFrank
cooked: <p>Hi Rane!</p>
<aside class="quote no-group" data-username="Rane90" data-post="1" data-topic="168137" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rane90/48/47118_2.png" class="avatar"> Rane90:</div>
<blockquote>
<p>I have two input tensors &lt;X, Y&gt;, with shapes: (batch_size, num_rows, num_columns). The loss I want to compute is the correlation between the respective &lt;X, Y&gt; columns for each one of the samples.</p>
</blockquote>
</aside>
<p>It’s relatively straightforward to implement a loop-free batch version of your<br>
<code>cov()</code> function.  Furthermore, <code>torch.linalg.inv()</code> (which <code>torch.inverse()</code><br>
is an alias for) computes matrix inverses on a batch basis if you pass in a<br>
tensor with more than two dimensions.  Together, these should give you the<br>
tools to implement a loop-free version of your <code>Corr</code> loss.</p>
<p>Here is an illustration of a batch-<code>cov()</code> and batch-<code>inv()</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.13.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; def cov(x, y):
...     x_bar = x - x.mean(axis=0)
...     y_bar = y - y.mean(axis=0)
...     N = x_bar.shape[0]
...     return (y_bar.T @ x_bar).T / (N - 1)
...
&gt;&gt;&gt; def batchCov (x, y):
...     x_bar = x - x.mean (1).unsqueeze (1)
...     y_bar = y - y.mean (1).unsqueeze (1)
...     N = x_bar.shape[1]
...     return (y_bar.transpose (1, 2) @ x_bar).transpose (1, 2) / (N - 1)   # @ will do batch mm
...
&gt;&gt;&gt; x = torch.randn (3, 5, 5)
&gt;&gt;&gt; y = torch.randn (3, 5, 5)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.allclose (batchCov (x, y)[2], cov (x[2], y[2]))
True
&gt;&gt;&gt;
&gt;&gt;&gt; torch.allclose (torch.linalg.inv (x)[2], torch.linalg.inv (x[2]))
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Rane90
cooked: <p>Thank you very much <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"> <a class="mention" href="/u/kfrank">@KFrank</a>,</p>
<p>I’ve used your <code>batchCov</code> function and modified my <code>Corr</code> to make it loop-free.</p>