google-research / nested-transformer

Nested Hierarchical Transformer https://arxiv.org/pdf/2105.12723.pdf

Nested Hierarchical Transformer Official Jax Implementation

NesT is a simple method, which aggregates nested local transformers on image blocks. The idea makes vision transformers attain better accuracy, data efficiency, and convergence on the ImageNet benchmark. NesT can be scaled to small datasets to match convnet accuracy.

This is not an officially supported Google product.

Pretrained Models and Results

Model	Accuracy	Checkpoint path
Nest-B	83.8	gs://gresearch/nest-checkpoints/nest-b_imagenet
Nest-S	83.3	gs://gresearch/nest-checkpoints/nest-s_imagenet
Nest-T	81.5	gs://gresearch/nest-checkpoints/nest-t_imagenet

Note: Accuracy is evaluated on the ImageNet2012 validation set.

Tensorbord.dev

See ImageNet training logs at Tensorboard.dev.

Colab

Colab is available for test: https://colab.sandbox.google.com/github/google-research/nested-transformer/blob/main/colab.ipynb

Pytorch re-implementation

The timm library has incorporated NesT and pre-trained models in Pytorch.

Instruction on Image Classification

Environment setup

virtualenv -p python3 --system-site-packages nestenv
source nestenv/bin/activate

pip install -r requirements.txt

Evaluate on ImageNet

At the first time, download ImageNet following tensorflow_datasets instruction from command lines. Optionally, download all pre-trained checkpoints

bash ./checkpoints/download_checkpoints.sh

Run the evaluation script to evaluate NesT-B.

python main.py --config configs/imagenet_nest.py --config.eval_only=True \
  --config.init_checkpoint="./checkpoints/nest-b_imagenet/ckpt.39" \
  --workdir="./checkpoints/nest-t_imagenet_eval"

Train on ImageNet

The default configuration trains NesT-B on TPUv2 8x8 with per device batch size 16.

python main.py --config configs/imagenet_nest.py --jax_backend_target=<TPU_IP_ADDRESS> --jax_xla_backend="tpu_driver" --workdir="./checkpoints/nest-b_imagenet"

Note: See jax/cloud_tpu_colab for info about TPU_IP_ADDRESS.

Train NesT-T on 8 GPUs.

python main.py --config configs/imagenet_nest_tiny.py --workdir="./checkpoints/nest-t_imagenet_8gpu"

The codebase does not support multi-node GPU training (>8 GPUs). The models reported in our paper is trained using TPU with 1024 total batch size.

Train on CIFAR

# Recommend to train on 2 GPUs. Training NesT-T can use 1 GPU.
CUDA_VISIBLE_DEVICES=0,1 python  main.py --config configs/cifar_nest.py --workdir="./checkpoints/nest_cifar"

Cite

@inproceedings{zhang2021aggregating,
  title={Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding},
  author={Zizhao Zhang and Han Zhang and Long Zhao and Ting Chen and and Sercan Ö. Arık and Tomas Pfister},
  booktitle={AAAI Conference on Artificial Intelligence (AAAI)},
  year={2022}
}

nested-transformer's Issues

About swintransformer ON cifar

could you share your patch size and window size when you run swin on CIFAR

Training hours & Imagenet accuracy

Hello, thanks for sharing your interesting work.

I was trying to reproduce the NesT-T ImageNet result in this link using TPUs.

Here are my TPU-v3 8 cores result (link) by using exactly the same hyperparameters in imagenet_nest_tiny.py

As you can see, it takes 63 hours for training while your result takes 21 hours.
How can I reduce training hours such as your result?
If this difference came from the data loading time, could you tell me the types of data storage that you used?
Right now, I'm using the google cloud storage bucket for data storage.

Furthermore, I can see the accuracy difference around 0.5% (81.0 v.s. 81.5).
Could you explain this difference?

Discrepancies vs Table A1 in paper

I noticed some possible discrepancies of the architecture parameters here vs those in table A1 of the paper

For ImageNet models, is it correct that:

1. The table should say h=[3,3,4]?
2. The order of the scale_hidden_dims in the table is inverted. That is, hierarchies 1, 2 and 3 should say [4d, 4h] × 2, 1 [2d, 2h] × 2, 4 [d, h] × k, 16?

Swin Transformer on CIFAR

Could you please let me know the hyperparameters you used to train Swin Transformer on CIFAR100?

Regarding GradCAT implementation

Hi
I'm interested to work with the Nest model, however I'm facing difficulty with the implementation of GradCAT. Could you please share the implementation?

This seems like it would be a great option for increasing context window in sequences. Have you tried that yet?

Just suggesting you try a similar architecture with sequences rather than images.

Can you make the Cifar training log public?

I would like to know the training details. Thanks a lot if you can help me.

Found it !!! Hope for Pytorch Implement.

Model Converge Problem

I am training on a medium-scale dataset that consists of 100,000 images. The learning rate and weight decay as the same as your config but still not working. Any opinion?

Regards,
Khawar Islam