Bottleneck in training due to double for loop

username: notprime
cooked: <p>Greetings from Italy!</p>
<p>So, I coded a custom loss function to train my autoencoder. The problem is the following:<br>
the function i’m using was coded by someone else, and, for my purposes, it needs to work on vectors.<br>
For instance, if I have a dataset <code>n x m</code> where <code>n</code> is the number of samples, and <code>m</code> is the number of features, I need the function to work on every possible combination of features: so I’ll pass to the function the features number 0 and 1, 0 and 2 and so on.<br>
So I did write a double for loop to see if the funtion would have worked inside torch. The problem is that the double for loop is being a huge bottleneck for my model, and the training is really slow.<br>
I tried to optimize the <strong>double for loop + function</strong> section, to vectorialize it, but I wans’t able to.<br>
This is the code:</p>
<pre><code class="lang-auto">
def rbf_dot(self, pattern1, pattern2, deg):
        size1 = pattern1.size() 
        size2 = pattern2.size()

        G = torch.sum(pattern1 * pattern1, axis = 1).reshape(size1[0], 1)
        H = torch.sum(pattern2 * pattern2, axis = 1).reshape(size2[0], 1)

        Q = torch.tile(G, (1, size2[0]))
        R = torch.tile(torch.transpose(H, 0, 1), (size1[0], 1))

        H = Q + R - 2 * torch.matmul(pattern1, torch.transpose(pattern2, 0, 1))

        H = torch.exp(-(H/2) / (deg**2))

        return H


    def function(self, X, Y):
        """
        X, Y : torch tensor of size [n, 1]
        n    : number of observations 
        """
        n = X.size(0)
        # ----- width of X -----
        Xmed = X

        G = torch.sum(Xmed * Xmed, axis = 1).reshape(n, 1)
        Q = torch.tile(G, (1, n))
        R = torch.tile(torch.transpose(G, 0, 1), (n, 1))
        dists = Q + R - 2 * torch.matmul(Xmed, torch.transpose(Xmed, 0, 1))
        dists = dists - torch.tril(dists)
        dists = dists.reshape(n**2, 1)
        width_x = torch.sqrt( 0.5 * torch.median(dists[dists&gt;0]))


        # ----- width of Y -----
        Ymed = Y

        G = torch.sum(Ymed * Ymed, axis = 1).reshape(n, 1)
        Q = torch.tile(G, (1, n))
        R = torch.tile(torch.transpose(G, 0, 1), (n, 1))

        dists = Q + R - 2 * torch.matmul(Ymed, torch.transpose(Ymed, 0, 1))
        dists = dists - torch.tril(dists)
        dists = dists.reshape(n**2, 1)
        width_y = torch.sqrt( 0.5 * torch.median(dists[dists&gt;0]))

        # ----- -----
        H = torch.eye(n, device = dists.device) - ( torch.ones(n,n, device = dists.device) / n ) 

        K = self.rbf_dot(X, X, width_x)
        L = self.rbf_dot(Y, Y, width_y)

        Kc = torch.matmul(torch.matmul(H, K), H)
        Lc = torch.matmul(torch.matmul(H, L), H)

        testStat = torch.sum(torch.transpose(Kc, 0, 1) * Lc) / n

        return testStat

    def double_for_loop_section(self, data):
        # n : rows,     observations
        # m : columns,  features
        # so that: a row represents all the features of a single user;
        #          a column represents the same feature for all users

        n, m = data.size()
        stat = torch.zeros(m, m, device = data.device)
        # self.device

        for i in range(m):
            for j in range(i, m):
                if i != j:
                    x = data[:, i]
                    y = data[:, j]
                    x_2D = x[:, None] 
                    y_2D = y[:, None]
                    stat[i, j] = self.function(x_2D, y_2D)
                    # matrix is symmetric
                    stat[j, i] = stat[i, j]
        return stat
</code></pre>
<p>The <code>double_for_loop</code> section is called inside the <code>loss</code> function method, passing the hidden layer output of the encoder.<br>
Could someone help me to vectorialize it?</p>

username: Zhaoyi-Yan
cooked: <p>Hi, give a more simplified version of your code, just a toy case that contains all the essential computation, maybe I can help.</p>

username: KFrank
cooked: <p>Hi Composite!</p>
<p>As an aside, your code is unhelpfully obtuse.</p>
<aside class="quote no-group quote-modified" data-username="notprime" data-post="1" data-topic="120759" 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/notprime/48/38089_2.png" class="avatar"> notprime:</div>
<blockquote>
<pre><code class="lang-auto">    def function(self, X, Y):
        ...
        Kc = torch.matmul(torch.matmul(H, K), H)
        Lc = torch.matmul(torch.matmul(H, L), H)

        testStat = torch.sum(torch.transpose(Kc, 0, 1) * Lc) / n
        ...
</code></pre>
</blockquote>
</aside>
<p>You loop over pairs of columns from <code>data</code> and then call <code>self.function()</code><br>
on that pair.</p>
<p>But up until the line in which you calculate <code>testStat</code>, <code>Kc</code> depends only<br>
on the first column in the pair and <code>Lc</code> only on the second.  So outside<br>
of your double loop you can precompute everything that goes into <code>Kc</code><br>
and <code>Lc</code>.</p>
<p>First do this precomputation to clear away all of the confusion.</p>
<p>Then, to “vectorize” the computation of <code>testStat</code> and placing in into the<br>
matrix <code>stat</code>, take a look at <code>torch.einsum()</code>.</p>
<p>Good luck!</p>
<p>K. Frank</p>

username: notprime
cooked: <p>Unofrtunately a simpler version of the code would be useless, all the operations you see are really important in computing what I need.<br>
Btw I’ve been able to vectorize it :P, thanks anyway!</p>

username: notprime
cooked: <p>This is exactly what I was doing.<br>
I didn’t write the code, I just translated it from numpy to torch.<br>
I also realized the code was utterly redundant!<br>
Btw, I’ve been able to vectorize it <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"><br>
I just need to vectorize one last thing which, for now, I assumed being constant <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"><br>
Thank you for the help, <strong>einsum</strong> was so cool to use <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>