Intuitive understanding of torch nn. Unfold

torch.nn.Unfold It's a batch Data in press C、Kernel_W、Kernel_H pack , Detailed explanation reference ：

PyTorch in torch.nn.functional.unfold Function usage details

This article is mainly about Unfold Back to tensor The middle part of the is restored to patches.

# -*- coding:utf-8 -*-
import cv2
import torch
import numpy as np

img1 = cv2.imread('../128128/1.png')
img2 = cv2.imread('../128128/2.png')

batch = torch.tensor([img1, img2]).permute(0,3,1,2)
print(batch.shape) # [2,3,128,128]

hor_block_num = 2
ver_block_num = 4

n,c,h,w = batch.shape

assert not h%ver_block_num
assert not w%hor_block_num

patch_h, patch_w = h//ver_block_num, w//hor_block_num

use_unfold = True
if (not use_unfold):
    #  Method 1 
    patches = batch.view(n, c, ver_block_num, patch_h, hor_block_num, patch_w)
    print(patches.shape) # [2,3,4,32,2,64]
    patches = patches.permute(0,2,4,3,5,1)
    print(patches.shape) # [2,4,2,32,64,3]
else:
    #  Method 2 
    split_block = torch.nn.Unfold(kernel_size=(patch_h, patch_w), stride=(patch_h, patch_w))
    patches = split_block(batch.float())
    print(patches.shape) # [2,6144,8]
    patches = patches.permute(0,2,1).view(n,ver_block_num,hor_block_num,-1,patch_h,patch_w)
    print(patches.shape) # [2,4,2,3,32,64]
    patches = patches.permute(0,1,2,4,5,3).byte()
    print(patches.shape) # [2,4,2,32,64,3]

#  preservation patch
for i in range(patches.shape[0]):
    img = patches[i]    
    for y in range(ver_block_num):
        for x in range(hor_block_num):
            patch_filename = '../128128/%d_y%d_x%d.png'%(i+1,y,x)
            patch = img[y,x].numpy()
            cv2.imwrite(patch_filename, patch)

The effect is as follows ：