[Pytorch] torch.nn.PixelShuffle/PixelUnshuffle 작동 방식

파이토치 라이브러리 중 하나인 pixelshuffle과 pixelunsuffle 함수의 작동 방식을 간단히 살펴보려한다.

Image Restoration 분야에서도 down sampling과 up sampling을 할 때 많이 활용하는 함수이다. 

말 그대로 픽셀들끼리 섞어서 그룹을 만들어 채널을 줄이면서 사이즈를 늘리거나(pixel shuffle) 반대로 픽셀들끼리 그룹을 만들고 채널 별로 펼쳐서 채널을 늘리면서 사이즈를 줄이는 방식이다(pixelunsuffle). 공식 설명은 아래와 같이 되어있다.

작동 방식은 정확히 안 나와있고 적용했을 때 feature 사이즈가 어떻게 변하는 지만 나와있는데 refer 되어있는 논문에 들어가서 관련 figure를 찾아보았다.

 위의 그림에서는 예시로 upscaling factor(위의 그림에서 r)가 3인 pixel shuffle을 보여주고 있다. channel 수는 1/(3²) 으로 줄어들면서 각 채널 별로 픽셀을 가져와 3×3 그룹으로 만들어 feature의 사이즈는 H와 W각각 3배로 증가한다. 위의 과정을 역으로 하면 pixel unshuffle이 된다. 

 일반적으로 down sampling을 한다고 하면 채널 수를 2배 늘리고 H,W를 각각 1/2배로 줄이면서 진행을 많이 하는데 위의 pixel unshuffle에서 downscale factor를 2로 설정하면 채널이 4배 늘고 H,W는 각각 1/2배로 준다. H,W는 원하는 만큼 줄였는데 채널이 원하는 것보다 2배 더 많이 늘어서 보통 1×1 convolution과 함께 적용해서 채널수를 절반 줄이면서 같이 사용한다.

 마찬가지로 up sampling도 기본적으로 채널 수를 1/2배로 줄이면서 H,W를 각각 2배로 늘리는데 pixel shuffle에서 upscale factor를 2로 설정하면 채널이 1/4배로 줄고 H,W가 각각 2배로 늘어나서 1×1 convolution으로 채널 수를 2배 늘려주면서 같이 사용한다.