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.