How to Unfold a tensor into sliding windows and then fold it back

username: Joshua_Clancy
cooked: <p>Hello!</p>
<p>So I have been trying to unfold an image tensor up into multiple sliding windows and then fold it back into an image. I have found multiple threads about this but none that have solved my problem. Right now I have successfully unfolded my images up into sliding windows like this:</p>
<pre><code class="lang-auto">im = torch.arange(0, 81).view(1,1,9,9)
im2 = torch.arange(0, 81).view(1,1,9,9)
x = torch.cat((im, im2), dim=0) # shape = [2,1,9,9]

patches = x.unfold(2,3,1)
print("unfold: ", patches.shape) # [2,1,7,9,3]

patches = patches.unfold(3,3,1)
print("unfold 2: ", patches.shape) # [2, 1, 7, 7, 3, 3]

patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()
print("after permute: ", patches.shape) #[2, 7, 7, 1, 3, 3]

patches = patches.view(patches.size()[0], -1, patches.size()[4], patches.size()[5])
print(patches.shape) # [2, 49, 3, 3]

</code></pre>
<p>At this point I have a bunch of sliding windows that I can work with… but before building a network around them I decided to try and make sure that I could fold them back. I have been working on this for a while now and seem to be stuck.</p>
<p>I have tried using fold in these ways:</p>
<pre><code class="lang-auto">fold = nn.Fold(output_size = (9,9), kernel_size = (3,3))
together = fold(patches)
print(together.shape)

fold = nn.Fold(output_size = 9, kernel_size = 3)
together = fold(patches)
print(together.shape)
</code></pre>
<p>But I keep running into problems. One big concern is that the documentation for fold seems wrong to me. It says: … warning::<br>
Currently, only 4-D output tensors (batched image-like tensors) are<br>
supported.<br>
But when I try and fold 4d tensors I get this error:</p>
<p>Input Error: Only 3D input Tensors are supported (got 4D)</p>
<p>Any help would be much appreciated! Thanks.</p>

username: meeros
cooked: <p>I confirm that I’ve seen the exact same issue.</p>

username: pfox
cooked: <p>I’m experiencing similar issues. (pytorch version 1.4.0)</p>
<p>EDIT: by looking at the source code <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py</a>, I can clearly see a mismatch between documentation and implementation:</p>
<pre><code class="lang-auto">def fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1):
    # type: (Tensor, BroadcastingList2[int], BroadcastingList2[int], BroadcastingList2[int], BroadcastingList2[int], BroadcastingList2[int]) -&gt; Tensor  # noqa
    r"""Combines an array of sliding local blocks into a large containing
    tensor.
    .. warning::
        Currently, only 4-D output tensors (batched image-like tensors) are
        supported.
    See :class:`torch.nn.Fold` for details
    """
    if not torch.jit.is_scripting():
        if type(input) is not Tensor and has_torch_function((input,)):
            return handle_torch_function(
                fold, (input,), input, output_size, kernel_size, dilation=dilation,
                padding=padding, stride=stride)
    if input.dim() == 3:
        msg = '{} must be int or 2-tuple for 3D input'
        assert_int_or_pair(output_size, 'output_size', msg)
        assert_int_or_pair(kernel_size, 'kernel_size', msg)
        assert_int_or_pair(dilation, 'dilation', msg)
        assert_int_or_pair(padding, 'padding', msg)
        assert_int_or_pair(stride, 'stride', msg)

        return torch._C._nn.col2im(input, _pair(output_size), _pair(kernel_size),
                                   _pair(dilation), _pair(padding), _pair(stride))
    else:
        raise NotImplementedError("Input Error: Only 3D input Tensors are supported (got {}D)".format(input.dim()))
</code></pre>

username: ptrblck
cooked: <p>The inline docs seem to be misleading. Could you please create an <a href="https://github.com/pytorch/pytorch/issues">issue</a> so that we can improve it, please?</p>

username: pfox
cooked: <aside class="quote no-group" data-username="pfox" data-post="3" data-topic="55890">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/a6a055/48.png" class="avatar"> pfox:</div>
<blockquote>
<p>Currently, only 4-D output tensors (batched image-like tensors) are supported</p>
</blockquote>
</aside>
<p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>, I created the issue <a href="https://github.com/pytorch/pytorch/issues/37063" rel="noopener nofollow ugc">#37063</a>. Let me know if it should have been better to issue it as a bug report, instead of as a doc issue.</p>

username: ptrblck
cooked: <p>I think the method does the “right” thing as far as I understand the method and think it’s a documentation issue.</p>
<p>Thanks for creating the issue!<br>
Would you be interested in creating a fix for it?</p>

username: pfox
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. Bout the fix, I think <a class="mention" href="/u/alband">@albanD</a> fixed this in pull request <a href="https://github.com/pytorch/pytorch/pull/37099" rel="nofollow noopener">#37099</a>.</p>
<p>With regards to <a class="mention" href="/u/joshua_clancy">@Joshua_Clancy</a>’s question above. I managed to solve the non-overlapping case, i.e., unfolding an image into non-overlapping windows and then putting them back together into the original shape:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as f
torch.manual_seed(0)

N, C, H, W = 2, 1, 9, 9
k = 3

x = torch.randn(N, C, H, W)
print(x.shape) # torch.Size([2, 1, 9, 9])

assert H % k == 0 and W % k == 0

patches = x.unfold(2, k, k).unfold(3, k, k)         
print(patches.shape) # torch.Size([2, 1, 3, 3, 3, 3]) - (N, C, H//k, W//k, k, k) 

patches = patches.reshape(N, C, H//k, W//k, k*k)
print(patches.shape) # torch.Size([2, 1, 3, 3, 9])    - (N, C, H//k, W//k, k*k)

patches = patches.permute(0, 1, 4, 2, 3)
print(patches.shape) # torch.Size([2, 1, 9, 3, 3])    - (N, C, k*k, H//k, W//k)

patches = patches.squeeze(1)
print(patches.shape) # torch.Size([2, 9, 3, 3])       - (N, k*k, H//k, W//k)

patches = patches.view(N, k*k, -1)
print(patches.shape) # torch.Size([2, 9, 9])          - (N, k*k, H//k * W//k)

folded = f.fold(patches, (H, W), kernel_size=k, stride=k)
print(folded.shape)  # torch.Size([2, 1, 9, 9])

print(torch.eq(x, folded).all()) # True
</code></pre>
<p>I’m finding it harder, though, to solve the “overlapping” case, which is the one <a class="mention" href="/u/joshua_clancy">@Joshua_Clancy</a> was actually asking about:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as f
torch.manual_seed(0)

C, H, W = 1, 9, 9
k = 3 # kernel size
s = 1 # stride

# Let's assume that our batch has only 1 sample, for simplicity
x  = torch.arange(0, H*W).view(1, C, H, W)
# im2 = torch.arange(0, H*W).view(1, C, H, W)
# x = torch.cat((im, im2), dim=0).type(torch.float32) # shape = [2, 1, 9, 9]

N = x.shape[0]

patches = x.unfold(2, k, s)
print("unfold:        ", patches.shape) # [2, 1, 7, 9, 3]

patches = patches.unfold(3, k, s)
print("unfold 2:      ", patches.shape) # [2, 1, 7, 7, 3, 3]

patches = patches.reshape(N, C, H-(k-1), W-(k-1), k*k)
print("reshape :      ", patches.shape) # [2, 1, 7, 7, 3*3]

patches = patches.permute(0, 1, 4, 2, 3)
print("after permute: ", patches.shape) # [2, 1, 3*3, 7, 7]

patches = patches.squeeze(1)
print("squeeze:       ", patches.shape) # [2, 3*3, 7, 7]

patches = patches.view(N, k*k, -1)
print("view:          ", patches.shape) # [2, 3*3, 7*7]

folded = f.fold(patches.type(torch.float32), (H, W), kernel_size=k, stride=s, dilation=1)
print("folded:          ", folded.shape) # [1, 1, 9, 9
</code></pre>
<p>If we now compare <code>x</code> and <code>folded</code>:</p>
<pre><code class="lang-python">&gt;&gt;&gt; x
tensor([[[[ 0,  1,  2,  3,  4,  5,  6,  7,  8],
          [ 9, 10, 11, 12, 13, 14, 15, 16, 17],
          [18, 19, 20, 21, 22, 23, 24, 25, 26],
          [27, 28, 29, 30, 31, 32, 33, 34, 35],
          [36, 37, 38, 39, 40, 41, 42, 43, 44],
          [45, 46, 47, 48, 49, 50, 51, 52, 53],
          [54, 55, 56, 57, 58, 59, 60, 61, 62],
          [63, 64, 65, 66, 67, 68, 69, 70, 71],
          [72, 73, 74, 75, 76, 77, 78, 79, 80]]]])

&gt;&gt;&gt; folded.type(torch.int32)
tensor([[[[  0,   2,   6,   9,  12,  15,  18,  14,   8],
          [ 18,  40,  66,  72,  78,  84,  90,  64,  34],
          [ 54, 114, 180, 189, 198, 207, 216, 150,  78],
          [ 81, 168, 261, 270, 279, 288, 297, 204, 105],
          [108, 222, 342, 351, 360, 369, 378, 258, 132],
          [135, 276, 423, 432, 441, 450, 459, 312, 159],
          [162, 330, 504, 513, 522, 531, 540, 366, 186],
          [126, 256, 390, 396, 402, 408, 414, 280, 142],
          [ 72, 146, 222, 225, 228, 231, 234, 158,  80]]]], dtype=torch.int32)
</code></pre>
<p>we can see that that <code>folded</code> has summed over overlapping elements, which is the expected behavior as explained in a <code>Note</code> in the docs for <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.Fold" rel="nofollow noopener">Fold</a>.</p>
<p>So to retrieve the initial tensor, I guess we could divide each element by the number of times the filter was “passed” over it. For instance, in our 9x9 case, the centered, inner 7x7 region should be divided by 9:</p>
<pre><code class="lang-python">folded[:, :, k-1:-(k-1), k-1:-(k-1)] = folded[:, :, k-1:-(k-1), k-1:-(k-1)] / 9
</code></pre>
<p>giving:</p>
<pre><code class="lang-python">&gt;&gt;&gt; x
tensor([[[[ 0,  1,  2,  3,  4,  5,  6,  7,  8],
          [ 9, 10, 11, 12, 13, 14, 15, 16, 17],
          [18, 19, 20, 21, 22, 23, 24, 25, 26],
          [27, 28, 29, 30, 31, 32, 33, 34, 35],
          [36, 37, 38, 39, 40, 41, 42, 43, 44],
          [45, 46, 47, 48, 49, 50, 51, 52, 53],
          [54, 55, 56, 57, 58, 59, 60, 61, 62],
          [63, 64, 65, 66, 67, 68, 69, 70, 71],
          [72, 73, 74, 75, 76, 77, 78, 79, 80]]]])
&gt;&gt;&gt; folded
tensor([[[[  0,   2,   6,   9,  12,  15,  18,  14,   8],
          [ 18,  40,  66,  72,  78,  84,  90,  64,  34],
          [ 54, 114,  20,  21,  22,  23,  24, 150,  78],
          [ 81, 168,  29,  30,  31,  32,  33, 204, 105],
          [108, 222,  38,  39,  40,  41,  42, 258, 132],
          [135, 276,  47,  48,  49,  50,  51, 312, 159],
          [162, 330,  56,  57,  58,  59,  60, 366, 186],
          [126, 256, 390, 396, 402, 408, 414, 280, 142],
          [ 72, 146, 222, 225, 228, 231, 234, 158,  80]]]], dtype=torch.int32)
</code></pre>
<p>Note, for instance, that the corner pixels <code>((0, 0), (0, 8), (8, 0) and (8, 8))</code> should be divided by 1, since they only appeared once in a filter pass. Therefore, they already correspond to the value of the initial tensor at such coordinates.</p>
<p>I am not sure this is a clean way to approach this “retrieval” of the initial tensor after an <code>unfold</code> operation. Would it be possible that <code>fold</code> didn’t <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.Fold" rel="nofollow noopener">“calculate each combined value in the resulting large tensor by summing all values from all containing blocks”</a>?</p>
<p>But then gain, maybe there’s a reason as to why <code>fold</code> should function this way and I am just missing the whole point <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thanks!</p>