Hybrid- Swin-Transformers 리뷰

Efficient large-scale image retrieval with deep feature orthogonality and Hybrid- Swin-Transformers

 

안녕하세요 이번 논문은 Nvidia 의 Christof Henkel가 2021 Google Landmark recognition과 retrieval challenge 에서 우승하고 해당 대회에서 사용한 모델에 대해 쓴 논문입니다. 사실 코드가 아직 공개되지 않아서 방법론을 정확히 이해하진 못했으나 논문이 짧고 저자가 Kaggle 랭킹 1위이신 분이기도 해서 리뷰해봤습니다.

1. Introduction

[Abstract]
A model with deep orthogonal fusion of local and global features (DOLG) using an EfficientNet backbone as well as a novel Hybrid-Swin-Transformer is discussed and details how to train both architectures efficiently using a step-wise approach and a sub-center arcface loss with dynamic margins are provided.

Dataset

원본 데이터셋은 200,000 가 넘는 클래스들을 포함한 오백만장 이상의 이미지들로 구성되어 있습니다. 데이터셋 사이즈가 클 뿐만 아니라 내부 클래스들은 long-tailed 분포를 띄고있고 noisy label도 포함되어 있다고 합니다.
For retrieval the task is to find similar images in a database with respect to a given query image and is evaluated using mean Average Precision

2. Modeling

Both track에 있어서 (retrieval, recognition) 내부 클래스의 variance가 크고 클래스 수가 많았기 때문에, image descriptor , high dimensional feature vector을 모두 학습할 수 있는 딥러닝 네트워크를 설계했다.

두 가지 모델 모두 구조적으로 EfficientNet 기반의 CNN 인코더와 sub-center arcface classification head을 이용하였다. 학습에 있어서 모든 이미지 사이즈들을 정사각형으로 조정했고 rotate, cutout augmentation 을 적용하였다. 옵티마이저로는 weight decay 를 적용한 Adam 을 사용하였으며 학습률과 배치 사이즈는 모델마다 상이하다. 스케줄러는 cosine annealing lr을 warm-up 1 에폭마다 사용하였다.

2.1 DOLG-EfficientNet with sub-center arcface

자체적으로 개선된 DOLG모델을 사용하였는데 첫 번째로, ImageNet 사전학습된 EfficeintNet 인코더를 사용하였다.
인코더의 세번째 블럭 위에 local branch를 추가하여 세 가지 사이즈의 dilated-convolution 을 통해 1024차원의 local feature을 추출하였다. 이후 피쳐 차원으로 concatenated 되고 spatial 2d 어텐션을 거치게 되어, 하단의 인코더 이후 동일하게 1024차원으로 GeM pooling 을 거친 global feature과 orthogonal fusion이 된다.
이후 average pooling 을 거쳐 [FC-BN-PReLU] 의 neck을 통해 512 차원의 final descriptor가 생성된다. 학습에 있어서 k=3인 sub-center arcface head 를 거쳐서 81,313개의 클래스들을 예측한다.

[Training Approach]

1. 작은 이미지 사이즈의 GLDv2c를 이용해 10에폭으로 학습
2. 이후 noisy 데이터가 많은 GLDv2x를 중간 이미지 사이즈로 30-40에폭 학습
3. GLDv2x를 이용하여 큰 이미지 사이즈로 적은 에폭으로 fine-tuning

 

2.2 Hybrid-Swin-Transformer with sub-center arcface

두 번째 아키텍처는 최근 발전한 VIT를 이용하는 방법이다. vision transformer을 CNN 인코더 뒤에 붙여 Hybrid-CNN 트랜스포머 모델을 만들었다고 할 수 있다. 앞의 CNN 파트가 local feature을 추출하는 부분이라 본다면 하단의 Transformer 파트는 추출된 local features를 그래프 신경망을 통해 단일 벡터로 통합한다고 해석할 수 있다.

특히 various scales을 사용가능한 Swin-Transformer 모델을 사용했다. 트랜스포머의 인풋으로 EfficientNet 의 output 을 flatten하여 positional embedding와 aggreegated information을 의미하는 virtual node가 추가되어 들어간다 (VIT의 patch embedding 과 동일한 부분이다. conceptually cls token과 virtual node가 같은 것 같다).

Swin Transformer 이후에는 마찬가지로 virtual token의 feature를 neck[FC-BN-PReLU]를 거쳐 final descriptor를 얻게 된다. 똑같이 sub-center arcface head를 거쳐 81,313 클래스를 예측.

EfficientNet 과 Swin-Transformer는 모두 사전학습된 이미지넷 가중치를 개별적으로 필요하기 때문에 학습시 overflow를 조심해야한다. 특히 mixed precision 사용시.

[Training Approach]

1. initialize transformer -neck-head 를 small image size 로 GLDv2c 10 에폭 학습
2. 기존 패치임베딩을 2블록 EfficientNet 인코더로 교체 후 기존 neck-head 부분을 freeze하고 새롭게 추가된 CNN을 meidum image size로 1에폭 학습.
3. unfreeze하고 전체모델을 GLDv2x로 30-40에폭 학습.
4. 사전학습된 EfficientNet을 CNN과 Swin 사이에 추가 후 전체 모델을 GLDv2x이용 large size로 적은 에폭으로 학습.

Firstly, initialize the transformer-neck-head part by training on a small im- age size using GLDv2c for 10 epochs. Next, exchangethe transformers orignal patch embedding with a 2-block EfficientNet encoder, freeze the previously transformer- neck-head part and train the newly added CNN part for one epoch on medium size images. thirdly, unfreeze and train the whole model for 30-40 epochs using GLDv2x. Finally, insert a further pretrained EfficientNet block between the CNN and Swin-Transformer and finetune the whole model for a few epochs using large images and GLDv2x.

→ 이부분이 좀 어려운데, 첫번째는 패치임베딩 교체 전이니 원본이미지로 vit 처럼 한다는듯. 두 번째는 neck-head 부분이 freeze되면 교체된 EfficientNet 부분이 학습이 안되는건가? → 이게 3번째 unfreeze할때 학습되는 것으로 이해. 4번의 further pretrained EfficientNet은 3번까지 진행한 모델인 것인가..

2.3 Submission ensemble

두 트랙 모두 우승 솔루션은 3개의 DOLG와 3개의 Hybrid-Swin-Transformer 모델들을 다양한 EfficientNet 백본모델과 인풋 이미지 사이즈를 이용하여 8개의 모델을 앙상블 한 것이다. (나머지 2개는 전체 데이터셋으로 학습한 plain EfficientNet 모델)

특이점은 작은 사이즈의 원본 이미지를 이용하는 몇몇 모델의 경우, 마지막 훈련 스텝에서 이미지 사이즈를 키우지 않고, 첫 번째 컨볼루션 레이어의 stride를 2 에서 1 로 설정하여 이득을 봤다.

 

3. Conclusion

구글 랜드마크대회에 참여는 안해봤지만 일단 데이터셋이 많은 recognition task에서 효과가 좋을 것 같습니다.
소타모델인 Swin-transformer 뿐만 아니라 CutMix,Mixup 등의 Augmentation이나 cosineannealing 스케쥴러 등 캐글에서 모델 성능을 높이는 많은 기법들이 쓰인 것 같습니다.
학습을 잘못한 탓이겠지만, 실제로 다른 데이터셋에 학습을 시킨 결과를 swin과 비교하니 그다지 좋지는 않았습니다ㅠ

공식 코드는 아닌 것 같은데 tensorflow로 짜여진 코드가 있어서 공유하겠습니다.
https://github.com/innat/EfficientNet-Hybrid-Swin-Transformer