Stuck in creating custom Pooling layer in Pytorch

quangduyhcmut: <p>Hello everyone,<br>
I am studying about quaternion convolutional neural network following this repo on github: <a href="https://github.com/Orkis-Research/Pytorch-Quaternion-Neural-Networks" class="inline-onebox" rel="noopener nofollow ugc">GitHub - Orkis-Research/Pytorch-Quaternion-Neural-Networks: This repository is an update to all previous repositories with implementations of various Quaternion-valued Neural Networks in PyTorch</a></p>
<p>The repo’s work is great but i want to implement a max amplitude pooling layer to utilize the quaternion network.</p>
<p>The basic idea follows this paper: <a href="https://ieeexplore.ieee.org/document/8632910" class="inline-onebox" rel="noopener nofollow ugc">Quaternion Convolutional Neural Network for Color Image Classification and Forensics | IEEE Journals &amp; Magazine | IEEE Xplore</a><br>
The pooling will take 4 input layer, compute the amplitude (length) then apply a max pooling. The torch.max function return pooled result and indices for max values.</p>
<p>My question is how to apply these indices to the input layer to get pooled results.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/1/21ae10990f217c4daf0a1dd6881a620f08bd5621.png" data-download-href="https://discuss.pytorch.org/uploads/default/21ae10990f217c4daf0a1dd6881a620f08bd5621" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/1/21ae10990f217c4daf0a1dd6881a620f08bd5621_2_690x471.png" alt="image" data-base62-sha1="4NWGHq392sfWf8EOaMIV50FJNTz" width="690" height="471" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/1/21ae10990f217c4daf0a1dd6881a620f08bd5621_2_690x471.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/1/21ae10990f217c4daf0a1dd6881a620f08bd5621_2_1035x706.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/1/21ae10990f217c4daf0a1dd6881a620f08bd5621.png 2x" data-dominant-color="F3EDE1"><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">1107×757 48.7 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></p>
<p>In my custom pooling layer, indices can be a list or a tensor with shape (n, c, h, w); input has shape (n, 4c, h, w) and amplitude has shape (n, c, h, w)</p>
<p>This is my forward function:</p>
<p>def forward(self, x):</p>
<pre><code>    # x = F.pad(x, self._padding(x), mode='constant')

    amplitude = get_modulus(x, vector_form=True)

    input_size = x.size()

    amp_size = amplitude.size()

    

    x_r = get_r(x)

    x_i = get_i(x)

    x_j = get_j(x)

    x_k = get_k(x)

    amp = amplitude.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])

    print("amp unfold size: ", amp.size())

    x_r = x_r.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])

    x_i = x_i.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])

    x_j = x_j.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])

    x_k = x_k.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])

    """

    tensor.unfold(axis, size, step)

    Returns a view of the original tensor which contains all slices of size size from self tensor in the dimension dimension.

    Step between two slices is given by step.

    """       

    amp = amp.contiguous().view(amp.size()[:4] + (-1,))

    print("amp size after view: ",amp.size())

    x_r = x_r.contiguous().view(x_r.size()[:4] + (-1,))

    x_i = x_i.contiguous().view(x_i.size()[:4] + (-1,))

    x_j = x_j.contiguous().view(x_j.size()[:4] + (-1,))

    x_k = x_k.contiguous().view(x_k.size()[:4] + (-1,))

    """

    tensor.contiguous() =&gt; create a copy of tensor

    """

    amp_max, indice = torch.max(amp, dim = -1)

    indice = torch.squeeze(indice).tolist()

    # print(indice)

    """

    initialize output channel

    """

    # x_r_out = torch.zeros_like(x_r.mean(4))

    # x_i_out = torch.zeros_like(x_r.mean(4))

    # x_j_out = torch.zeros_like(x_r.mean(4))

    # x_k_out = torch.zeros_like(x_r.mean(4))

    """

    apply max amplitude position =&gt; max amplitude pooling

    """

    # x_r_out

    # x_i_out

    # x_j_out

    # x_k_out

    print('size of input: ', input_size)

    print('size of amplitude: ', amp_size)

    print('size of pooling: ', amp_max.size())

    

    return torch.cat((x_r_out, x_i_out, x_j_out, x_k_out), dim = 1)
</code></pre>

tom: <p>I’d imagine that it is something along the lines of</p>
<pre><code class="lang-python">qs = torch.randn(1, 4, 6, 6) # batch 1, 4 channels (1 quaternion) and 6x6 matrix, bs = 1
bs, _, h, w = qs.shape
amp = torch.norm(qs, dim=1, keepdim=True)
amp_uf = amp.unfold(3, 2, 2).unfold(2, 2, 2).reshape(bs, 1, h // 2, w // 2, 2 * 2)
qs_uf = qs.unfold(3, 2, 2).unfold(2, 2, 2).reshape(bs, 4, h // 2, w // 2, 2 * 2)
maxamp = amp_uf.argmax(-1, keepdim=True).expand(-1, 4, -1, -1, -1)
qs_pooled = qs_uf.gather(-1, maxamp).view(bs, 4, h // 2, w // 2)
</code></pre>
<p>(The gather gets goes over the dimensions except the last and then picks the value at the maxamp index in the last dimension.)</p>
<p>Best regards</p>
<p>Thomas</p>

quangduyhcmut: <p>That’s amazing! Thank you so much for a clever solution!</p>

SAAD_SAOUD_Lyes: <p>Thanks much, Thomas.<br>
It is really amazing. I wonder if you can help to come back from qs_pooled to qs (I mean upsampling or MaxUnpool2d PyTorch equivalent)<br>
Regards<br>
Lyes</p>