Need help understanding Conv2d and fold, unfold

ChrisLiu2: <p>I’m working on a cnn that directly processes each patch. After reading the documentation on fold and unfold, my understanding is that I can first apply convolution on an arbitrary <code>[b, c, h, w]</code> input named <code>A</code>, with some parameters for stride, dilation and padding. Let’s notate the output as shape <code>[b, c, h1, w1]</code>, named <code>B</code>.</p>
<p>My understanding for how fold and unfold works is as follows: If I were to unfold the input <code>A</code>, I would get something of the shape <code>[b, H, L]</code>. Then I can apply some transformation to <code>H</code> dimension, and then use unfold with <code>B</code>’s shape for the <code>output_shape</code> parameter, and taking the same parameters like stride and dilation.</p>
<p>However, it doesn’t look like it works this way. If my understanding is correct, the following example should execute without a problem. But it doesn’t run. Can someone point out what I missed?</p>
<pre><code class="lang-python">A = torch.randn(4, 32, 224, 224)
out_shape = nn.Conv2d(32, 32, kernel_size=3, dilation=2, stride=2, padding=0)(A).shape[-2:]
windows = torch.nn.functional.unfold(A, kernel_size=3, dilation=2, stride=2, padding=0)
A_ = torch.nn.functional.fold(windows, output_size=out_shape, kernel_size=3, dilation=2, stride=2, padding=0)
</code></pre>
<p>This would complain that fold given output size=110 would require a size of L=2809, while the windows variable has an L=12100 which is 110**2. Then I tried to fold again with output_size=224, which doesn’t complain anymore but gives me back a shape of <code>[4, 32, 224, 224]</code>.</p>

eqy: <p>Take a look at this documentation for the non-functional version:</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.Fold.html#torch.nn.Fold" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.Fold.html#torch.nn.Fold</a></p>
<p>In particular to:</p>
<pre><code class="lang-auto">fold(unfold(input)) == divisor * input
</code></pre>
<p>I believe <code>Fold</code> would restore the original input shape rather than yield the output shape of the corresponding convolution.</p>

ChrisLiu2: <p>I see, now that you mention it the doc does mention that it restores the original size. Do you know of a way to preserve the behavior of parameters such as dilation and stride, so that the effect of unfold and fold is almost identical to that of a Conv2d module?</p>

eqy: <p>I would think about your problem from the perspective that some sort of reduction needs to happen in order to transform the Unfold output to the equivalent of a conv. This example may not be the exact operation you are trying to implement but it illustrates some basics:</p>
<pre><code class="lang-auto">import torch
from torch import nn

A = torch.ones(4, 32, 224, 224)
conv = nn.Conv2d(32, 1, kernel_size=3, dilation=2, stride=2,padding=0, bias=False)
conv.weight = torch.nn.Parameter(torch.ones_like(conv.weight))
out = conv(A)
print(out.shape)
windows = torch.nn.functional.unfold(A, kernel_size=3, dilation=2, stride=2, padding=0)
print(windows.shape)
windows = windows.reshape(4, 1, 32*9, out.shape[-2], out.shape[-1])
windows_reduced = torch.sum(windows, axis=2)
print(windows_reduced.shape)
print(torch.allclose(out, windows_reduced))
</code></pre>
<pre><code class="lang-auto">torch.Size([4, 1, 110, 110])
torch.Size([4, 288, 12100])
torch.Size([4, 1, 110, 110])
True
</code></pre>
<p>Note that unfold will produce an “output” equivalent to that of a convolution with just a single channel.</p>

ChrisLiu2: <p>Thank you for the detailed reply. I do see this solution working and I’ll go with something along this line.</p>