[83] Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs (RepLKNet)

[dhkim0225]

paper
code - MegEngine
code - pytorch

개인적으로 좋아하는 MEGVII 의 Rep~ style 연구
31x31 large kernel 을 활용하여 좋은 성능을 이끌어낸다.

ERF 결과

RepLKNet

Custom Kernel for DW conv

FFT 를 이용하면 DW conv 를 상당히 가속화 시킬 수 있다.
https://arxiv.org/abs/1312.5851

FFT 의 특징 2가지.
Spatial domain의 Convolution --> Frequency domain의 element-wise multiplication
Spatial domain의 element-wise multiplication --> Frequency domain의 convolution

즉, 31x31 의 커다란 conv 를 frequency 로 옮겨서 element-wise multiplication 을 수행해 버린다.

하지만 이 방법보다는 block-wise implicit gemm algorithmn 이 더 좋은 방법이더라.
implicit gemm algorithm 은 아래 CUTLASS 글에 잘 설명되어 있다.
https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
https://github.com/NVIDIA/cutlass/blob/master/media/docs/implicit_gemm_convolution.md

참고로, 기존 GEMM 은 다음과 같다.

이렇게 만든 custom kernel 을 이용해 속도 측정을 해 보니, 빠르더라.

shortcut

shortcut 을 붙여주는게 확실히 Large-Kernel 학습에서 도움이 되더라.
MobileNetV2 에서 shortcut 을 붙여보고 빼봤는데, 다음과 같은 결과가 나왔다.

shortcut은 ViT 에서도 잘 먹힌다.

Re-parameterization

shortcut 이 있는 형태로 학습이 끝나고 나면 RepVGG 스타일로 kernel 을 합쳐 버린다.

recap 을 위해 RepVGG Figure 를 하나 가져왔다.
BN 까지 Fuse 해준 뒤, 각 inception style kernel 를 하나로 합쳐준다.

fuse code 는 다음과 같다.
https://github.com/DingXiaoH/RepLKNet-pytorch/blob/c61441ba800b6452ba2a54826c479279226852d4/replknet.py#L107-L114

재밌는 건, identity shortcut 보다, 3x3 conv 와 같이 small kernel 을 합치면 더 성능이 좋은 것이다.
MobileNet V2 에서 shortcut 쪽 kernel 크기를 더 늘리면 다음과 같은 성능이 나온다.

Large Conv boost downstream tasks

바로 위 table 성능을 보면 ImageNet 성능은 약간만 좋아졌는데, CityScape 성능이 눈에 띄게 좋아지는 걸 볼 수 있다.

Large kernel is useful even on small feature maps (e.g. 7x7)

last stage kernel 크기도 다양하게 변화시켜 봤는데, 큰 커널이 더 최종 성능이 좋더라.

저자들은 inductive bias 가 줄어들어서 (strict 하게 translational equivariance 를 강조하지 않아서) 좋은 것이라고 주장한다.

Architecture

최종 모델은 다음과 같다.
그림에는 보이지 않지만, 5x5 를 reparametrization 으로 함께 넣어줬다고 한다.
즉, (skip conn), (bn - 1x1 - 31x31 - 1x1), (bn - 1x1 - 5x5 - 1x1) 3개의 path 가 구성되는 형태

Results

‡ 는 22K pretriain 결과
다이아몬드는 MegData 73M 추가 사용 결과.

ImageNet

Swin-L 부터는 RepLKNet 이 좀 딸리는 느낌.
저자들도 limitation 에 해당 내용을 적어 놨다.

Cityscape (Semantic Segmentation)

ADE20K (Semantic Segmentation)

COCO (Object Detection)

Large-Kernel Models are More Similar to Human in Shape Bias

참 재밌는 결과다.
swin 은 확실히 texture 를 더 중시하는 편이다.
RepLKNet은 사람처럼 shape을 중시함.