ImageNet-training

Pytorch ImageNet training codes with various tricks, lr schedulers, distributed training, mixed precision training, DALI dataloader etc. We hope this repo can help ImageNet experiments in NAS researches.

Train

CUDA_VISIBLE_DEVICES=0 python -u train.py --train_root /path/to/imagenet/train_set --val_root /path/to/imagenet/val_set --train_list /path/to/imagenet/train_list --val_list /path/to/imagenet/val_list

Please refer to train_example.sh for more details.

Test

CUDA_VISIBLE_DEVICES=0 python -u test.py --val_root /path/to/imagenet/val_set --val_list /path/to/imagenet/val_list --weights /path/to/pretrained_weights

Please refer to test_example.sh for more details.

Model Profiling

Please refer to profile_example.py for more details.

Tested on

Python == 3.7.6
pytorch == 1.5.1
torchvision == 0.6.1
nvidia.dali == 0.22.0
cuDNN == 7.6.5
apex from this link

License

This repo is released under the MIT license. Please see the LICENSE file for more information.