Unfold a tensor

Johnson_Mark: <p>Hello all, I have a tensor size of <code>BxCxHxW</code>. I want to unfold the tensor with a kernel size of <code>K</code> into non-overlapped patches. Do we have any equation to compute the stride and padding for the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html" rel="noopener nofollow ugc">unfold</a> function, such that the patches can be used to fold the original tensor <code>BxCxHxW</code> by fold function?</p>
<p>For example, a tensor size of <code>16x32x56x56</code> undolds with size of <code>k=6</code>, which should I use stride and padding value?</p>

ptrblck: <p>For non-overlapping patches the stride should equal the kernel size.<br>
For a small example have a look at <a href="https://discuss.pytorch.org/t/creating-nonoverlapping-patches-from-3d-data-and-reshape-them-back-to-the-image/51210/6">this post</a>.<br>
Based on your input shape and kernel size, you would need to pad the input tensor to properly reshape the patches back.</p>

Johnson_Mark: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> this is my example, however, it does not work for recovering back the original tensor when h and w are not divisible to k. This is my code</p>
<pre><code class="lang-auto">import torch
from torch.nn import functional as F

def unfold_tensor (x, step_c, step_h, step_w):
    kc, kh, kw = step_c, step_h, step_w  # kernel size
    dc, dh, dw = step_c, step_h, step_w  # stride
    pad_c, pad_h, pad_w = x.size(1)%kc // 2,  x.size(2)%kh // 2, x.size(3)%kw // 2
    x = F.pad(x, ( pad_h, pad_h, pad_w, pad_w, pad_c, pad_c))
    patches = x.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
    unfold_shape = patches.size()    
    patches = patches.reshape(-1,unfold_shape[1]*unfold_shape[2]*unfold_shape[3], unfold_shape[4]*unfold_shape[5]*unfold_shape[6])
    return patches, unfold_shape

def fold_tensor (x, shape_x, shape_orginal):
    x = x.reshape(-1,shape_x[1], shape_x[2], shape_x[3], shape_x[4], shape_x[5], shape_x[6])
    x = x.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
    #Fold
    output_c = shape_x[1] * shape_x[4]
    output_h = shape_x[2] * shape_x[5]
    output_w = shape_x[3] * shape_x[6]
    x = x.view(1, output_c, output_h, output_w)
    return x

b,c,h,w = 1, 8, 28, 28
x = torch.randint(10, (b,c,h,w))
print (x.shape)
shape_orginal = x.size()
patches, shape_patches = unfold_tensor (x, c, 3, 3)
print (patches.shape)
fold_patches = fold_tensor(patches, shape_patches, shape_orginal)
print (fold_patches.shape)
print((x == fold_patches).all())
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">torch.Size([1, 8, 28, 28])
torch.Size([1, 81, 72])
torch.Size([1, 8, 27, 27])
Traceback (most recent call last):
  File "test_unfold.py", line 32, in &lt;module&gt;
    print((x == fold_patches).all())
</code></pre>

ptrblck: <p>The used padding is too naive in my example and you might want to use e.g. <code>divmod</code> to calculate the padding size:</p>
<pre><code class="lang-python">def unfold_tensor (x, step_c, step_h, step_w):
    kc, kh, kw = step_c, step_h, step_w  # kernel size
    dc, dh, dw = step_c, step_h, step_w  # stride
    
    nc, remainder = np.divmod(x.size(1), kc)
    nc += bool(remainder)
    
    nh, remainder = np.divmod(x.size(2), kh)
    nh += bool(remainder)
    
    nw, remainder = np.divmod(x.size(3), kw)
    nw += bool(remainder)    
    
    pad_c, pad_h, pad_w = nc*kc - x.size(1),  nh*kh - x.size(2), nw*kw - x.size(3)
    x = F.pad(x, ( 0, pad_h, 0, pad_w, 0, pad_c))
    patches = x.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw)
    unfold_shape = patches.size()    
    patches = patches.reshape(-1,unfold_shape[1]*unfold_shape[2]*unfold_shape[3], unfold_shape[4]*unfold_shape[5]*unfold_shape[6])
    return patches, unfold_shape

def fold_tensor (x, shape_x, shape_orginal):
    x = x.reshape(-1,shape_x[1], shape_x[2], shape_x[3], shape_x[4], shape_x[5], shape_x[6])
    x = x.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
    #Fold
    output_c = shape_x[1] * shape_x[4]
    output_h = shape_x[2] * shape_x[5]
    output_w = shape_x[3] * shape_x[6]
    x = x.view(1, output_c, output_h, output_w)
    return x

b,c,h,w = 1, 8, 28, 28
x = torch.randint(10, (b,c,h,w))
print(x.shape)
shape_original = x.size()
patches, shape_patches = unfold_tensor (x, c, 3, 3)
print(patches.shape)
fold_patches = fold_tensor(patches, shape_patches, shape_orginal)
print(fold_patches.shape)
fold_patches = fold_patches[:, :shape_original[1], :shape_original[2], :shape_original[3]]
print(fold_patches.shape)
print((x == fold_patches).all())
&gt; tensor(True)
</code></pre>
<p>Currently the padding is only applied to one side and you can try to split it so that both sides will be padded.</p>