Implementation of momentum^2 teacher
- All experiments are done with python3.6, torch==1.5.0; torchvision==0.6.0
Data Preparation
Prepare the ImageNet data in ${root_of_your_clone}/data/imagenet_train
, ${root_of_your_clone}/data/imagenet_val
. Since we have an internal platform(storage) to read imagenet, I have not tried the local mode. You may need to do some modification in momentum_teacher/data/dataset.py
to support the local mode.
Training
Before training, ensure the path (namely ${root_of_clone}
) is added in your PYTHONPATH, e.g.
export PYTHONPATH=$PYTHONPATH:${root_of_clone}
To do unsupervised pre-training of a ResNet-50 model on ImageNet in an 8-gpu machine, run:
- using
-d
to specify gpu_id for training, e.g.,-d 0-7
- using
-b
to specify batch_size, e.g.,-b 256
- using
--experiment-name
to specify the output folder, and the training log & models will be dumped to ‘./outputs/${experiment-name}’ - using
-f
to specify