Awesome
Accelerating Vision Transformer Training via a Patch Sampling Schedule (PSS)
We introduce the notion of a Patch Sampling Schedule (PSS), that varies the number of Vision Transformer (ViT) patches used per batch during training. Since all patches are not equally important for most vision objectives (e.g., classification), we argue that less important patches can be used in fewer training iterations, leading to shorter training time with minimal impact on performance. Additionally, we observe that training with a PSS makes a ViT more robust to a wider patch sampling range during inference. This allows for a fine-grained, dynamic trade-off between throughput and accuracy during inference. We evaluate using PSSs on ViTs for ImageNet both trained from scratch and pre-trained using a reconstruction loss function.
Model Checkpoints and Logs
Below are checkpoints and training/validation logs for baseline models and models + PSS. All training times are for a single machine with 4 A100 GPUs.
| Codebase | Model | Checkpoint | Logs | Pretrained? | Training Time | Accuracy |
|---|---|---|---|---|---|---|
| SimMIM | ViT-B | checkpoint | logs | yes | 24.6h | 83.7% |
| SimMIM | ViT-B + PSS | checkpoint | logs | yes | 17.0h | 83.5% |
| DeiT | DeiT-S-224 | checkpoint | logs | no | 29.9h | 80.1% |
| DeiT | DeiT-S-224 + PSS | checkpoint | logs | no | 27.9h | 80.4% |
| DeiT | DeiT-S-384 + PSS | checkpoint | logs | no | 109.0h | 82.7% |
Training and Inference results
Please refer to the simmim/ or deit/ directories for more information on how to train and validate PSS models.