This is the PyTorch source code for the Mean Teacher paper. The code runs on Python 3. Install the dependencies and prepare the datasets with the following commands:
pip install numpy scipy pandas pytorch tqdm matplotlib
pip install git+ssh://git@github.com/pytorch/vision@c31c3d7e0e68e871d2128c8b731698ed3b11b119
The code expects to find the data in specific directories inside the data-local directory. You can prepare the CIFAR-10 with this command:
./data-local/bin/prepare_cifar10.sh
You can prepare the ImageNet using these instructions (Section "Download the ImageNet dataset"). The mean teacher code expects to find the ImageNet data at data-local/images/ilsvrc2012/
.
To train on CIFAR-10, run e.g.:
python main.py \
--dataset cifar10 \
--labels data-local/labels/cifar10/1000_balanced_labels/00.txt \
--arch cifar_shakeshake26 \
--consistency 100.0 \
--consistency-rampup 5 \
--labeled-batch-size 62 \
--epochs 180 \
--lr-rampdown-epochs 210
Use python main.py --help
to see other command line arguments.
To reproduce the CIFAR-10 ResNet results of the paper run python -m experiments.cifar10_test
using 4 GPUs.
To reproduce the ImageNet results of the paper run python -m experiments.imagenet_valid
using 10 GPUs.