Speed up custom function and model using GPU

username: jose
cooked: <p>Hello everyone,</p>
<p>I am working on implementing some custom function (and its corresponding module) that will be part of the sequential execution in a layer. When running it I get a very bad performance in terms of speed of training. I even commented the computations and left only the loops and the calls to the custom function, and even then it is still super slow. It could be that the looping is the main cause. This seems unlikely, however, since the size of the tensor is [20, 1, 28, 28] and [20, 32, 12, 12], corresponding to the two layers. The most interesting parts of the code are the following:</p>
<pre><code class="lang-auto">class FrankWolfe_MinCutSum_Canonical(Function):
    @staticmethod
    def in_bounds(ctx, coord):
        (x, y) = coord
        return 0 &lt;= x &lt; ctx.width and 0 &lt;= y &lt; ctx.height

    @staticmethod
    def neighbors(ctx, coord):
        (x, y) = coord
        return list(filter(lambda p: FrankWolfe_MinCutSum_Canonical.in_bounds(ctx, p), [(x+1, y), (x, y-1), (x-1, y), (x, y+1)]))

    @staticmethod
    def relaxed_taylor_closed_form_solution(ctx, w):
        u_star = ctx.beta
        for x in range(ctx.width):
            for y in range(ctx.height):
                i = (x, y)
                u_star[i] -= ctx.alpha*sum([(1 if w[i] &gt; w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, i)])
        return u_star

    @staticmethod
    def forward(ctx, beta, alpha=1, tol=1e-6, max_iter=20):
        # print(beta)
        ctx.height, ctx.width = list(beta.shape)
        ctx.alpha, ctx.beta = alpha, beta
        v = beta
        # for t in range(max_iter):
        #     w = v.exp().div(v.exp().add(1))
        #     u_star = FrankWolfe_MinCutSum_Canonical.relaxed_taylor_closed_form_solution(ctx, w)
        #     gamma_t = 2 / (t + 2)
        #     v_1 = v
        #     v = (1 - gamma_t)*v_1 + gamma_t*u_star
        #     if torch.norm(v - v_1) &lt; tol:
        #         break
        ctx.save_for_backward(v.clone())
        return v

    @staticmethod
    def backward(ctx, grad_output):
        v, = ctx.saved_variables
        # return v * grad_output
        return grad_output

class LateralInteractions(nn.Module):
    def __init__(self):
        super(LateralInteractions, self).__init__()

    def forward(self, x):
        out = x
        batch_size, channels, _, _ = list(x.shape)
        for b in range(batch_size):
            for c in range(channels):
                # continue
                # out[b, c] = FrankWolfe_MinCutSum_Canonical.apply(Variable(x[b, c], requires_grad=True))
                out[b, c] = FrankWolfe_MinCutSum_Canonical.apply(x[b, c])
        return out
</code></pre>
<p>I left in the comments so that you can see that right now the function is basically passing the input as output without changing anything. Then the network looks as follows:</p>
<pre><code class="lang-auto">class ConvNetLat(nn.Module):
    def __init__(self, n = 10):
        super(ConvNetLat, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=7, stride=2, padding=1),
            LateralInteractions())
        self.layer2 = nn.Sequential(
            nn.Conv2d(32, n, kernel_size=7, stride=1, padding=1),
            LateralInteractions(),
            nn.AdaptiveAvgPool2d(1))
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, x):
        # [20, 1, 28, 28]
        out = self.layer1(x)
        # [20, 32, 12, 12]
        out = self.layer2(out)
        # [20, 10, 1, 1]
        out = out.reshape(out.size(0), -1)
        out = self.log_softmax(out)
        return out
</code></pre>
<p>What could I do to improve the training performance? This is literally taking hours to train, as is (with commented code and all).</p>
<p>Any help will be greatly appreciated.</p>

username: iidsample
cooked: <p>Most probably it’s the for loops. For better performance it’s always better do vector operations.<br>
so the line -<br>
u_star[i] -= ctx.alpha*sum([(1 if w[i] &gt; w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, i)])</p>
<p>if you can write it as a vector operation without a for loop it should be much faster.</p>

username: jose
cooked: <p>Done. Changed the following:</p>
<pre><code class="lang-auto">    def relaxed_taylor_closed_form_solution(ctx, w):
        u_star = ctx.beta
        u_star -= (torch.cuda if torch.cuda.is_available() else torch).FloatTensor([ \
                    [ctx.alpha*sum([(1 if w[(x, y)] &gt; w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, (x, y))]) \
                    for y in range(ctx.height)] for x in range(ctx.width)])
        return u_star
</code></pre>
<p>and</p>
<pre><code class="lang-auto">class LateralInteractions(nn.Module):
    def __init__(self):
        super(LateralInteractions, self).__init__()

    def forward(self, x):
        # out = x
        batch_size, channels, _, _ = list(x.shape)
        out = torch.stack([ \
                torch.stack([FrankWolfe_MinCutSum_Canonical.apply(Variable(x[b, c], requires_grad=True)) \
                for c in range(channels)]) for b in range(batch_size)])
        # out[b, c] = FrankWolfe_MinCutSum_Canonical.apply(x[b, c])
        return out
</code></pre>
<p>Still super slow though. Any other recommendations? Maybe consider running in parallel could improve enough?</p>

username: iidsample
cooked: <p>You are still running a for loop when you are doing for loop in the stack, i think you can do it one go for c and b and then do another stack. Those for loops seems to be the bottleneck.</p>

username: jose
cooked: <p>Is this the correct way?</p>
<pre><code class="lang-auto">    def relaxed_taylor_closed_form_solution(ctx, w):
        u_star = ctx.beta
        u_star -= (torch.cuda if torch.cuda.is_available() else torch).FloatTensor([ \
                    ctx.alpha*sum([(1 if w[(x, y)] &gt; w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, (x, y))]) \
                    for x in range(ctx.width) for y in range(ctx.height)]).reshape(ctx.width, ctx.height)
        return u_star
</code></pre>
<pre><code class="lang-auto">class LateralInteractions(nn.Module):
    def __init__(self):
        super(LateralInteractions, self).__init__()

    def forward(self, x):
        # out = x
        batch_size, channels, height, width = list(x.shape)
        out = torch.stack([FrankWolfe_MinCutSum_Canonical.apply(Variable(x[b, c], requires_grad=True)) \
                for b in range(batch_size) for c in range(channels)]).reshape(batch_size, channels, height, width)
        # out[b, c] = FrankWolfe_MinCutSum_Canonical.apply(x[b, c])
        return out
</code></pre>
<p>Everything else is as in the original post.</p>

username: jose
cooked: <p>After digging more into it I have come to the conclusion that the comparison of scalar tensors is taking a huge time. Any way how to make this faster?</p>
<pre><code class="lang-auto">        return ctx.beta - torch.FloatTensor([ \
                ctx.alpha*np.sum([(1 if w[(x, y)] &gt; w[j] else -1) for j in FrankWolfe_MinCutSum_Canonical.neighbors(ctx, (x, y))]) \
                for x in range(ctx.height) \
                for y in range(ctx.width)]).reshape(ctx.width, ctx.height).to(device)
</code></pre>

username: iidsample
cooked: <p>One of the things is that why are you not doing all the ops on GPU, might give you some speed up.<br>
And again try to remove the for loop and vectorize the operations. Will require some tricks like you will have to unroll the matrix like for each ctx.width create a corresponding row or column of all ctx.weight instead of looping over.</p>

username: jose
cooked: <p>I changed the function into this:</p>
<pre><code class="lang-auto">    def relaxed_taylor_closed_form_solution(ctx, w):
        w_shift_left = torch.cat((w[:, 1:], w[:, -1].reshape(-1, 1)), 1)
        w_shift_right = torch.cat((w[:, 0].reshape(-1, 1), w[:, :-1]), 1)
        w_shift_up = torch.cat((w[1:, :], w[-1, :].reshape(1, -1)), 0)
        w_shift_down = torch.cat((w[0, :].reshape(1, -1), w[:-1, :]), 0)
        interactions = ctx.alpha*(torch.sign(w - w_shift_left) + torch.sign(w - w_shift_right) + torch.sign(w - w_shift_up) + torch.sign(w - w_shift_down))
        return ctx.beta - interactions
</code></pre>
<p>The execution was still too slow, so I started investigating on parallelization. Currently, I implemented this, which works good for CPU but crashes with GPU:</p>
<pre><code class="lang-auto">from joblib import Parallel, delayed
import multiprocessing

class LateralInteractions(nn.Module):
    def __init__(self):
        super(LateralInteractions, self).__init__()

    def forward(self, x):
        batch_size, channels, height, width = x.shape
        num_cores = multiprocessing.cpu_count()
        results = Parallel(n_jobs=num_cores, backend="threading")(delayed(FrankWolfe_MinCutSum_Canonical.apply)(x[b, c]) for b in range(batch_size) for c in range(channels))
        out = torch.stack(results).reshape(batch_size, channels, height, width)
        return out
</code></pre>
<p>To give you a general idea, the execution is faster for parallelized CPU than for GPU without parallelization. It would be ideal to make that part of the code to be called in parallel, since the calls are completely independent and only need to be synchronized before passing to the next layer.</p>

username: qi_cao
cooked: <p>Thanks a lot! This parallelization method really helps me!</p>