[FEATURE] ViT with average pooling like ResNet

[Phuoc-Hoan-Le]

Is there any attempt on training ViT models with global average pooling like resnet models? If yes, what are the exact hyperparameters used to get the best performance when training on imagenet1k from scratch

[rwightman]

@CharlesLeeeee there are already some there, but discussions for these sorts of q please.

[Phuoc-Hoan-Le]

@rwightman I see "vit_base_patch16_rpn_224" which uses average pooling and res-post norm residuals, but I have no idea what learning rate; warmup-epochs; batch size; and data augmentation techniques you used for that model.

I just used the exact DeiT training recipe and I got NaN.

I cannot find any discussions related to training "vit_base_patch16_rpn_224". If you know one could you point me the link? Thanks!

[rwightman]

Most of these use global pooling, but they're relative position https://github.com/rwightman/pytorch-image-models/blob/main/timm/models/vision_transformer_relpos.py

They don't really train much differently than token models though. Most of the newer weights for these are trained with a swin type recipe, which isn't far from deit. Deit should not blow up, it's possible you didn't adapt the LR to your batch size or didn't enable grad clipping... I have hparams somewhere in a cloud bucket but have to find them...

[Phuoc-Hoan-Le]

I never had to use gradient clipping to prevent it from going to NaN, especially as in DeiT 1 paper (https://arxiv.org/pdf/2012.12877.pdf). I find that using warmup-epochs=20 solves the problem most of the time.

[Phuoc-Hoan-Le]

So far I am able to train a DeiT model with average pooling successfully by increasing the warmup-epochs from 5 to 20.

However, when using a res-post-norm structure with avg pooling if it keeps on getting NaN.

Are you using repeated-aug to get these results? I know that the swin model doesn't use repeated-aug

[Doraemonzzz]

Hello, by changing the warmup to 20 and using gap, is it possible to replicate the results of the deit paper?

[rwightman]

vit-b16_rpn-args.yaml.txt

hparams from the rpn b16 attached, it was trained on TPU v4-8 so it's 4x256 global batch size, but it's using --lr-base so it autoscales the LR ...