Speeding up custom convolutions for multiple CPUS

Tom_Ginsberg: <p>Hi I’m trying to implement a convolution that operates on a slighly different image representation.</p>
<p>This representation can be though of as taking an 2D image with pixel values <code>x_i</code><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb9cb6b6494dd15fb7ecf24833adbaf13248aef7.png" alt="image" data-base62-sha1="qLH2LjT9C08YgOUgsg1DV9om5a7" width="170" height="80"><br>
And adding some additional terms for each pixel<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/3/c/3cdaa46ad98cd0a72b104a556cf694a9b4247d61.png" alt="image" data-base62-sha1="8Gl5oO3V3da9UCMIYjxpGUq3O8x" width="176" height="215"><br>
I see that this feels similar as a multi channel image, but the import difference is the inner most vectors should be treated as a single object and added element wise. For example a more complex image like<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/b/1b34cb5be164a8af394058afe2d9bf17921915dd.png" data-download-href="https://discuss.pytorch.org/uploads/default/1b34cb5be164a8af394058afe2d9bf17921915dd" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1b34cb5be164a8af394058afe2d9bf17921915dd_2_169x250.png" alt="image" data-base62-sha1="3SG0bZwB4HTKxbA1R6LUlwYeXXn" width="169" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1b34cb5be164a8af394058afe2d9bf17921915dd_2_169x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1b34cb5be164a8af394058afe2d9bf17921915dd_2_253x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/b/1b34cb5be164a8af394058afe2d9bf17921915dd_2_338x500.png 2x" data-dominant-color="EBEBEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">564×832 51.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
When convolved with the kernel<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/3/7/371b0109d6ab37b4db90cdb657fa04f8f7136bf2.png" alt="image" data-base62-sha1="7Ru7eYo7cFgnyJVpNpNZ0Nqn0Ma" width="72" height="58"><br>
Would output this mess<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/2/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198.png" data-download-href="https://discuss.pytorch.org/uploads/default/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/2/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198_2_517x213.png" alt="image" data-base62-sha1="e6eF2M2KJOUP5W7s1GJ793c071u" width="517" height="213" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/2/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198_2_517x213.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/2/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198_2_775x319.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/2/62d2e8a3b26a4446f4ddb7250f1d32fcfc7e1198_2_1034x426.png 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1532×632 40 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
The other fun twist is bias must only be added to the first element of every pixel vector.</p>
<p>I have a simple implementation ignoring strides and padding that looks like this</p>
<pre><code class="lang-auto">    def forward(self, img):
        k = self.kernel_size
        # dimension of each image channel (assumes n x n)
        n = img.shape[-2]

        # &lt;5&gt; stack all output channels
        return torch.stack(
            # &lt;4&gt; take all channels, stack and sum
            [torch.stack(
                # &lt;3&gt; takes all rows and stack
                [torch.stack(
                    # &lt;2&gt; take inner products across a row and stack
                    [torch.stack(
                        # &lt;1&gt; take inner product with kernel at top left corner pixel (l, m)
                        [torch.stack([kernel[i, j] * img_channel[l + i, m + j] for i in range(k) for j in
                                      range(k)]).sum(0).index_add(0, tensor([0]), bias)
                         # &lt;1&gt;
                         for m in range(n - k)])
                        # &lt;2&gt;
                        for l in range(n - k)])
                    # &lt;3&gt;
                    for img_channel in img]).sum(0)
             # &lt;4&gt;
             for kernel, bias in zip(self.weight, self.bias)])
        # &lt;5&gt;
</code></pre>
<p>I figured since this is really so similar to regular convolution exept that pixels are now vectors and added elementwise someone might have an idea for how I could use some pytorch builtins more effectively. Since I will be running this on a 14 core machine this nasty thing of list comprehensions will really hurt my performance.</p>
<p>Any ideas are much appreciated. Thanks</p>

albanD: <p>Hi,</p>
<p>I can take a stab at this.<br>
Could you make a self contained script with your implementation, random inputs of the exact shape you’re interested in (including batch potentially and specifying which dimension is which with your math above) and the expected output?</p>

Tom_Ginsberg: <p>Sure just give me a minute.</p>

Tom_Ginsberg: <p><a href="https://gist.github.com/tomginsberg/172733c91e545648a8321d1e149047a2" rel="nofollow noopener">Here’s</a> a self contained script with a short description and several test cases. The help is really appreciated. Thx</p>

albanD: <p>Hey,</p>
<p>Here is an implementation that should be significantly faster and match what you want. Because of the very different low precision of your ground truth, the test trips, but comparing the two forward passes lead to very small error (same up to 5 places):</p>
<pre><code class="lang-python">class Conv2D(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(Conv2D, self).__init__()
        # These are ones only for testing
        self.weight = Parameter(torch.ones(out_channels, in_channels, kernel_size, kernel_size), requires_grad=False)
        self.bias = Parameter(torch.ones(out_channels), requires_grad=False)
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding

    def forward(self, img):
        # img is (num_channels x n x n x num_errors)
        # outer dimension of each image channel (assumes n x n)
        num_channels = img.size(0)
        n = img.shape[-2]
        # kernel is square
        k = self.kernel_size
        # padding is the same in both H and W
        p = self.padding

        # Consider the last dimension as a batch
        img = img.permute(-1, 0, 1, 2)
        # Use regular conv
        out = F.conv2d(img, self.weight, None, self.stride, self.padding)
        # Add custom biases only on the first entry
        out[0] += self.bias.view(-1, 1, 1) * num_channels
        # Reshape to the original input
        out = out.permute(1, 2, 3, 0)
        return out
</code></pre>
<p>Note that you could do similar stuff with a batch of input by doing a similar permute as <code>(0, -1, 1, 2, 3)</code>. Then collapse the first dimension in a common batch. After the conv, uncollapse the two dimensions and permute back.</p>