Home

Awesome

Accepted to ICCV 2021 !!

Rethinking Spatial Dimensions of Vision Transformers

Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, Seong Joon Oh | Paper

NAVER AI LAB

<img src="teaser.png" width="100%" title="" alt="teaser"></img>

News

Abstract

Vision Transformer (ViT) extends the application range of transformers from language processing to computer vision tasks as being an alternative architecture against the existing convolutional neural networks (CNN). Since the transformer-based architecture has been innovative for computer vision modeling, the design convention towards an effective architecture has been less studied yet. From the successful design principles of CNN, we investigate the role of the spatial dimension conversion and its effectiveness on the transformer-based architecture. We particularly attend the dimension reduction principle of CNNs; as the depth increases, a conventional CNN increases channel dimension and decreases spatial dimensions. We empirically show that such a spatial dimension reduction is beneficial to a transformer architecture as well, and propose a novel Pooling-based Vision Transformer (PiT) upon the original ViT model. We show that PiT achieves the improved model capability and generalization performance against ViT. Throughout the extensive experiments, we further show PiT outperforms the baseline on several tasks such as image classification, object detection and robustness evaluation.

Model performance

We compared performance of PiT with DeiT models in various training settings. Throughput (imgs/sec) values are measured in a machine with single V100 gpu with 128 batche size.

NetworkFLOPs# paramsimgs/secVanilla+CutMix+DeiT+Distill
DeiT-Ti1.3 G5.7 M256468.768.572.274.5
PiT-Ti0.71 G4.9 M303071.372.673.074.6
PiT-XS1.4 G10.6 M212872.476.878.179.1
DeiT-S4.6 G22.1 M98068.776.579.881.2
PiT-S2.9 G23.5 M126673.379.080.981.9
DeiT-B17.6 G86.6 M30369.375.381.883.4
PiT-B12.5 G73.8 M34876.179.982.084.0

Use PiT models with timm repo

Install timm>=0.4.7 using:

pip install git+https://github.com/rwightman/pytorch-image-models.git

Create PiT models

import torch
import timm

model = timm.create_model('pit_s_224', pretrained=True)
print(model(torch.randn(1, 3, 224, 224)))

Pretrained weights

Model nameFLOPsaccuracyweights
pit_ti0.71 G73.0link
pit_xs1.4 G78.1link
pit_s2.9 G80.9link
pit_b12.5 G82.0link
pit_ti_distilled0.71 G74.6link
pit_xs_distilled1.4 G79.1link
pit_s_distilled2.9 G81.9link
pit_b_distilled12.5 G84.0link

Dependancies

Our implementations are tested on following libraries with Python 3.6.9 and CUDA 10.1.

torch: 1.7.1
torchvision: 0.8.2
timm: 0.3.4
einops: 0.3.0

Install other dependencies using the following command.

pip install -r requirements.txt

How to use models

You can build PiT models directly

import torch
import pit

model = pit.pit_s(pretrained=False)
model.load_state_dict(torch.load('./weights/pit_s_809.pth'))
print(model(torch.randn(1, 3, 224, 224)))

Or using timm function

import torch
import timm
import pit

model = timm.create_model('pit_s', pretrained=False)
model.load_state_dict(torch.load('./weights/pit_s_809.pth'))
print(model(torch.randn(1, 3, 224, 224)))

To use models trained with distillation, you should use _distilled model and weights.

import torch
import pit

model = pit.pit_s_distilled(pretrained=False)
model.load_state_dict(torch.load('./weights/pit_s_distill_819.pth'))
print(model(torch.randn(1, 3, 224, 224)))

License

Copyright 2021-present NAVER Corp.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Citation

@inproceedings{heo2021pit,
    title={Rethinking Spatial Dimensions of Vision Transformers},
    author={Byeongho Heo and Sangdoo Yun and Dongyoon Han and Sanghyuk Chun and Junsuk Choe and Seong Joon Oh},
    booktitle = {International Conference on Computer Vision (ICCV)},
    year={2021},
}