[Pytorch] torch.nn.functional.unfold / fold

이번에는 이미지를 공간(Spatial) 도메인에서 Patch 단위로 분해하고 다시 조립하는 과정에 대해 정리해보려 한다. (관련 문서는 다음 링크 참조: https://docs.pytorch.org/docs/stable/generated/torch.nn.Unfold.html)

딥러닝, 특히 Vision Transformer(ViT)나 커스텀 Convolution 연산을 구현할 때 이미지를 작은 Patch로 잘라내야 하는 경우가 많다. 이때 사용되는 것이 torch.nn.functional.unfold와 그 역변환인 fold이다.

먼저 torch.nn.functional.unfold이다. 이 함수는 이미지에서 슬라이딩 윈도우를 적용해 \(k \times k\) 크기의 블록들을 추출하고, 이를 평탄화하여 반환해준다. 주요 Parameter로는 kernel_size, stride, padding 등이 있다. 이미지를 \(C \times H \times W\) 형태라고 가정하고, \(3 \times 3\) 커널로 패치를 자른다고 하면:

unfolded = torch.nn.functional.unfold(image, kernel_size=3, stride=3)

여기서 중요한 건 Output의 Shape이다. 변환된 unfolded 텐서는 (C * k * k, L) 형태가 된다.

dim 0: 패치 내의 모든 픽셀 값 (Channel \(\times\) kernel_h \(\times\) kernel_w)

dim 1 (L): 추출된 총 패치의 개수

즉 아래와 같은 이미지가 있다고 할때 (3*128*128)

이를 kernel size를 32, stride 32로 해서 나누게 되면 아래와 같이 나눠질 것이다.

이 패치들이 각각 flatten되어 합쳐지면서 실제 output인 unfolded 텐서의 shape은 (3*32*32, 16)이 되게 된다.

이를 다시 원래대로 합쳐주고 싶으면 아래와 같이 output size를 지정해주고 나눌때 사용했던 kernel size와 stride를 넣어서 다시 fold함수를 사용해주면 된다.

H_out, W_out = 128, 128
folded_image = torch.nn.functional.fold(unfolded, output_size=(H_out, W_out), kernel_size=32, stride=32)

예시에서는 kernel과 stride 사이즈를 같게 해서 겹치는 부분이 없게 해서 보였지만 실제로는 겹치는 부분이 생기게 자르는 경우가 많다.