Unfold + Linear + Fold != Conv2d

pbelcak: <p>Hello,</p>
<p>This has appeared both on the forums and in PyTorch issues before (<a href="https://github.com/pytorch/pytorch/issues/47990" rel="noopener nofollow ugc">this one is still open</a>). Nevertheless, I still find that when training a “simulated” Conv2d using <code>Unfold, Transpose, Linear, Transpose, Fold</code>, the gradients are different to using just the “equivalent” Conv2d.</p>
<p>For a reproducible example, I provide a “simulated” VGGxx implementation and loss curves from corresponding runs (all parameters being equal).</p>
<p>To get the simulated VGG, copy the contents of the torchvision implementation and modify <code>make_features</code> as follows:</p>
<pre><code class="lang-auto">def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -&gt; nn.Sequential:
    layers: List[nn.Module] = []
    in_channels = 3
    pooling_depth: int = 0
    for v in cfg:
        if v == "M":
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            pooling_depth += 1
        else:
            v = cast(int, v)
            w, h = 224 // (2**pooling_depth), 224 // (2**pooling_depth)
            layers += [ nn.Unfold(kernel_size=3, padding=1, stride=1) ]
            layers += [ Transpose(-1, -2) ]
            layers += [ nn.Linear(in_channels * 3*3, v, bias=True) ]
            layers += [ nn.ReLU(inplace=True) ]
            layers += [ Transpose(-1, -2) ]
            # layers += [ Reshape(-1, v, w, h) ]
            layers += [ nn.Fold(output_size=(w, h), kernel_size=1, stride=1) ]
            
            in_channels = v
    return nn.Sequential(*layers)
</code></pre>
<p>The training harness used is the one of <a href="https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg" rel="noopener nofollow ugc"><code>torchvision/classification</code></a>, and I follow the recipe given in the README.</p>
<p>Here are the losses for the first few epochs of VGG11 reference (blue) and simulated (yellow):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/5/55f277346dd4ff64489130d533ee86a9be428f55.png" data-download-href="https://discuss.pytorch.org/uploads/default/55f277346dd4ff64489130d533ee86a9be428f55" title="Screenshot 2024-05-02 at 19.50.38"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/5/55f277346dd4ff64489130d533ee86a9be428f55.png" alt="Screenshot 2024-05-02 at 19.50.38" data-base62-sha1="cgk42L6TvoPoTp7Bu5sq0sZGhp3" width="690" height="497" data-dominant-color="F3F5FA"><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">Screenshot 2024-05-02 at 19.50.38</span><span class="informations">874×630 22.8 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>Note that I also tried training both in double precision and it didn’t make a difference.</p>
<p>Am I missing something, or is there genuinely a bug in <code>Unfold/Fold</code> backward pass?</p>

ptrblck: <p><a href="https://discuss.pytorch.org/t/how-to-implement-a-convolutional-layer/68211/7">This post</a> as well as the docs show an example code snippet showing the same results between a manual unfold approach and the corresponding conv layer. Did you check each layer separately making sure the outputs are equal up to floating point precision?</p>

pbelcak: <p>Yes, on a forward pass the outputs of the two approaches agree. It is the backward pass where they don’t seem to agree, and the issue I referenced gives a snippet where they show whether they agree or not is a bit random but does improve with increasing the float precision.</p>
<p>I seem to be experiencing the same issue when trying to modify the reference implementations of ResNet and AlexNet.</p>

pbelcak: <p>Moreover, if I use Conv2d for a few layers and Unfold+Linear for others, the model begins to converge, but much slower than if it were full Conv2d</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/3/b36f3b9e1c5683b7708ee9754af0909aaaa6b8e5.png" data-download-href="https://discuss.pytorch.org/uploads/default/b36f3b9e1c5683b7708ee9754af0909aaaa6b8e5" title="Screenshot 2024-05-02 at 20.48.21"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/3/b36f3b9e1c5683b7708ee9754af0909aaaa6b8e5.png" alt="Screenshot 2024-05-02 at 20.48.21" data-base62-sha1="pBlMw9QvhalgSlfw6dFu3l7rwnr" width="690" height="497" data-dominant-color="F1F1F8"><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">Screenshot 2024-05-02 at 20.48.21</span><span class="informations">874×630 22.9 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>

ptrblck: <p>Could you post a standalone layer showing the differences in the backward pass, please?</p>