nn.Conv2d with custom offset operation

decadenza: <p>I need to implement a 2D convolution that operates on xyz coordinates.</p>
<p>My data is made by ordered points, with shape HxW and each point has 3 coordinates.<br>
Ideally I can use <code>nn.Conv2D</code> and slide over exactly like an image.</p>
<p>However, before the first convolution with kernel size 3x3 I need to subtract the central point coordinates from each of the neighbours. In my understanding, this is to build a “local” feature.</p>
<p>Question is: how can I efficiently write this custom <code>nn.Conv2d</code>? Can I do it with <code>nn.Module</code>?</p>
<p>I checked <a href="https://discuss.pytorch.org/t/custom-nn-conv2d/62068">here</a> but I wasn’t able to find an example of unfold. Also, not sure if it would be of reasonable performance.<br>
Apologies if it is a silly question.</p>

Sami_Hassan: <p>Hello, No question is silly, please don’t worry these forums are very helpful for people like us.</p>
<p>As i understand from your question you need to do a step before convolution layer right. i think all you need is hooks for your task. Please check this link if it is similar to your problem. <a href="https://pytorch.org/tutorials/beginner/former_torchies/nnft_tutorial.html" rel="noopener nofollow ugc">[1]</a></p>

decadenza: <p>Thanks.</p>
<p>Option 1. Assuming the following layer:</p>
<pre><code class="lang-auto">class MyConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm3d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.25),
        )

    def forward(self, x):
        return self.conv(x)
</code></pre>
<p>Should I do something like:</p>
<pre><code class="lang-auto">def mycustomfunction(self, input, output):
    # do my stuff
    # I need to access each local block *just before* convolution
    # Should I use unfold here? How?
    
layer = MyConv(3,64)
layer.conv[0].register_forward_hook(mycustomfunction)
</code></pre>
<p>Option 2, would be something like:</p>
<pre><code class="lang-auto">class MyConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        
        self.conv = nn.Sequential(
            
            ### Not sure how to use nn.Unfold here!

            nn.BatchNorm3d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.25),
        )
    
    def forward(self, x):
        return self.conv(x)
</code></pre>
<p>I simply need to subtract the central item from the all local block before the convolution.</p>

decadenza: <p>In the end, adapting the code from <a href="https://github.com/pytorch/pytorch/issues/47990" rel="noopener nofollow ugc">this issue</a>, I was able to do the following:</p>
<pre><code class="lang-auto">class OffsetConv2d(nn.Conv2d):
    """
    2d convolution with offset of central features
    
    See Also
    --------
    https://github.com/pytorch/pytorch/issues/47990
    """
    
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0, bias=True):
        super().__init__(in_channels, out_channels, kernel_size, stride, padding, bias)
        self.k = kernel_size
        self.stride = stride
        self.padding = padding
        self.conv = nn.Parameter(torch.ones(out_channels, in_channels, kernel_size, kernel_size))
        
    def forward(self, x):
        k = self.k
        stride = self.stride
        h_in, w_in = x.shape[2], x.shape[3]

        padding = self.padding
        batch_size = x.shape[0]

        h_out = (h_in + 2 * padding - (k - 1) - 1) / stride + 1
        w_out = (w_in + 2 * padding - (k - 1) - 1) / stride + 1
        h_out, w_out = int(h_out), int(w_out)
        
        inp_unf = F.unfold(x, (k, k), padding=padding)
        
        
        # DO MY STUFF HERE
        # Offset each local block with its center
        
        out_unf = inp_unf.transpose(1, 2).matmul(self.conv.view(self.conv.size(0), -1).t()).transpose(1, 2)
        out = F.fold(out_unf, (h_out, w_out), (1, 1))
                        
        return out
</code></pre>
<p>I compared the above code with the traditional nn.Conv2d and I didn’t find any significant performance difference. So I will go ahead with this.<br>
I hope that it is right and the above code will be useful to someone else in the same situation.</p>