Improving Function Performance

legoh: <p>Hi folks, I’m seeking some help with my function below. I’d like to get some opinions on if there’s a way to improve the inference performance of my function.</p>
<p>Here’s some context of what I want to do. Beliefs is <code>batch x n x m</code> matrix, where each <code>batch</code> is basically a factor graph with <code>n</code> nodes (rows of the matrix) and <code>m</code> factors (columns of the matrix), where the graph is bipartite (<code>n</code> nodes on one side, <code>m</code> factors on the other, edges only from <code>n</code> to <code>m</code> and vice versa). From this, I calculate messages known as <code>fac2var</code> and <code>var2fac</code> which are factor to variable messages and variable to factor messages respectively. For factors, the message it receives is the 2nd lowest incoming message from the variable nodes (<code>n</code> nodes). As you can see, this is done via the <code>for-loop</code> in the message propagation. That section is extremely slow at the moment, as there’s an iterative <code>for-loop</code> that is going through.</p>
<p>Basically, I loop through all <code>n</code> nodes, and mask out the message it would be sending for each factor. Then, for every factor, I will find the 2nd lowest message it receives. This is done in such a fashion because for a factor <code>m</code>, the message it sends to <code>n</code> will be the 2nd lowest message it receives, NOT including the message it received from node <code>n</code>. Hence, I need to mask out the messages accordingly.</p>
<p>I was wondering if this can be done in a <code>torch.scatter</code> or <code>torch.gather</code> way, which could potentially speed it up a fair bit? I’m open to suggestions too!</p>
<p>Thanks a lot for the help.</p>
<pre><code class="lang-auto">    def constraint_bp(beliefs, adj, bp_iter):
        # For batch-level implementation, the beliefs is arranged as [batch_size, max_n, max_m]

        # We can construct a few matrices to handle the problem. We have 3 matrices, belief matrix (b x nx1),
        # adjacency matrix (b x n x m) and the factor-&gt;var matrix (b x n x m),
        # where n is the number of variable nodes, and m the number of factors
        batch_size, n, m = adj.size()
        fac2var = torch.zeros((batch_size, n, m), dtype=torch.float, device=beliefs.device)
        # Add a mask which inflates values by 99999 so that disconnected nodes are not considered in the message aggregation
        mask = adj.clone().detach()
        mask[mask == 0] = 99999
        mask[mask == 1] = 0

        lamb = 0.2
        # Find all variable to factor messages
        var2fac = torch.multiply(beliefs, adj)
        # BP algorithm

        # Note that for an updated var2fac message, we can simply take the beliefs subtract away by the
        # previous fac2var message for that variable. The local beliefs were
        # initialized with d(e), which is the local factor for the variable. Hence, the total beliefs of the node
        # is given by d(e) + all incoming messages (which are essentially damped updates). Thus, we can always just take
        # the message on the edge out. Also note that if we were to convert this problem from MIN to MAX,
        # we need to be careful of how we treat d(e). Right now, there's a general assumption
        # that all updates are POSITIVE, because of the nature of d(e).

        for rounds in range(bp_iter):
            if rounds != 0:
                var2fac = torch.multiply(beliefs, adj) - fac2var

            # Compile message for every factor, we can use the adjacency matrix to mask away the msg
            # To calculate the message of column_j (fac_j) to row_i (var_i), we take the var2fac messages,
            # mask away row_i, take the min/max of column_j
            fac2var_tilde = torch.zeros((batch_size, n, m), dtype=torch.float, device=beliefs.device)
            masked_var2fac = var2fac + mask
            for i in range(n):
                neighbors = torch.cat((masked_var2fac[:, :i, :], masked_var2fac[:, (i + 1):, :]), dim=1)
                fac2var_tilde[:, i, :] = torch.topk(neighbors, k=2, dim=1, largest=False)[0][:, 1, :]
            fac2var_tilde = torch.multiply(fac2var_tilde, adj) * -1.0

            eps = fac2var_tilde - fac2var
            fac2var = fac2var + lamb * eps
            beliefs = beliefs + torch.sum(eps, dim=2).unsqueeze(2)

        return beliefs
</code></pre>