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.pyis the rough version from scratch. - The
main_single.pyis the old version that is only able to train on a single gpu card. - The
main.pyis 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 1We provide three three modes,
train,validateandtest. The argment-gindicates the gpu ids. Add-r 3if 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 |