This repo is a pytorch implementation of the paper. We aim to reprogram the pretrained imagenet models for MNIST classification. Following the code, you can easily add more datasets and other pretrained imagenet models for more experiments.
- The
main_raw.py
is the rough version from scratch. - The
main_single.py
is the old version that is only able to train on a single gpu card. - The
main.py
is the final version that supports multi-gpu training.
-
Download pretrained imagenet models.
./download.py
-
The MNIST dataset will be automatically downloaded after running the scripts. Then, the structure of this repo should be as follows.
├── datasets │ └── mnist │ ├── processed │ │ ├── test.pt │ │ └── training.pt │ └── raw │ ├── t10k-images-idx3-ubyte.gz │ ├── t10k-labels-idx1-ubyte.gz │ ├── train-images-idx3-ubyte.gz │ └── train-labels-idx1-ubyte.gz ├── models │ ├── resnet50-19c8e357.pth │ └── vgg16-397923af.pth ├── train_log │ └── W_001.pt ├── README.md ├── download.py ├── main.py ├── ....
-
So directly run the following scripts to train the model.
python main.py -m train -g 0 1
We provide three three modes,
train
,validate
andtest
. The argment-g
indicates the gpu ids. Add-r 3
if you want to continue training your model from the third epoch.
Simply running the following command.
python main.py -m validate -g 0 -r 30
The performance (test accuracy) of ImageNet model reprogrammed on other tasks. More experiments are welcomed by pull requests.
ImageNet Model | MNIST |
---|---|
Resnet50 | 0.94 |
VGG16 | 0.94 |