This repo contains the official Pytorch reimplementation of the paper "NetAdapt: Platform-Aware Neural Network Adaptation for Mobile Applications" [paper] [project]. The results in the paper were generated by the Tensorflow implementation from Google AI.
The code base is tested with the following setting:
- Python 3.7.0
- CUDA 10.0
- Pytorch 1.2.0
- torchvision 0.4.0
- numpy 1.17.0
- scipy 1.3.1
First clone the repo in the directory you want to work:
git clone https://github.com/denru01/netadapt.git
cd netadapt
In the following context, we assume you are at the repo root.
If the versions of Python and CUDA are the same as yours, you can download the python packages using:
pip install -r requirements.txt
To verify the downloaded code base is correct, please run either
sh scripts/unittest.sh
or
sh scripts/unittest_helloworld.sh
sh scripts/unittest_alexnet.sh
sh scripts/unittest_mobilenet.sh
If it is correct, you should not see any FAIL.
In order to apply NetAdapt, run:
python master.py [-h] [-gp GPUS [GPUS ...]] [-re] [-im INIT_MODEL_PATH]
[-mi MAX_ITERS] [-lr FINETUNE_LR] [-bu BUDGET]
[-bur BUDGET_RATIO] [-rt RESOURCE_TYPE]
[-ir INIT_RESOURCE_REDUCTION]
[-irr INIT_RESOURCE_REDUCTION_RATIO]
[-rd RESOURCE_REDUCTION_DECAY]
[-st SHORT_TERM_FINE_TUNE_ITERATION] [-lt LOOKUP_TABLE_PATH]
[-dp DATASET_PATH] [-a ARCH] [-si SAVE_INTERVAL]
working_folder input_data_shape input_data_shape
input_data_shape
-
working_folder
: Root folder where models, related files and history information are saved. You can see how models are pruned progressively inworking_folder/master/history.txt
. -
input_data_shape
: Input data shape (C, H, W) (default: 3 224 224). If you want to apply NetAdapt to different tasks, you might need to change data shape. -
-h, --help
: Show this help message and exit. -
-gp GPUS [GPUS ...], --gpus GPUS [GPUS ...]
: Indices of available gpus (default: 0). -
-re, --resume
: Resume from previous iteration. In order to resume, specify--resume
and specifyworking_folder
as the one you want to resume. The resumed arguments will overwrite the arguments provided here. For example, if you want to simplify a model by pruning and finetuning for 30 iterations (underworking_folder
), however, your program terminated after 20 iterations. Then you can use--resume
to restore and continue for the last 10 iterations. -
-im INIT_MODEL_PATH, --init_model_path INIT_MODEL_PATH
: Path to pretrained model. -
-mi MAX_ITERS, --max_iters MAX_ITERS
: Maximum iteration of removing filters and short-term fine-tune (default: 10). -
-lr FINETUNE_LR, --finetune_lr FINETUNE_LR
: Short-term fine-tune learning rate (default: 0.001). -
-bu BUDGET, --budget BUDGET
: Resource constraint. If resource <budget
, the process is terminated. -
-bur BUDGET_RATIO, --budget_ratio BUDGET_RATIO
: If--budget
is not specified,buget
=budget_ratio
*(pretrained model resource) (default: 0.25). -
-rt RESOURCE_TYPE, --resource_type RESOURCE_TYPE
: Resource constraint type (default: FLOPS). We currently supportFLOPS
,WEIGHTS
, andLATENCY
(devicecuda:0
). If you want to add other resource types, please modifydef compute_resource(...)
innetwork_util
python files (e.g.network_utils/network_utils_alexnet
). -
-ir INIT_RESOURCE_REDUCTION, --init_resource_reduction INIT_RESOURCE_REDUCTION
: For each iteration, target resource = current resource -init_resource_reduction
*(resource_reduction_decay
**(iteration-1)). -
-irr INIT_RESOURCE_REDUCTION_RATIO, --init_resource_reduction_ratio INIT_RESOURCE_REDUCTION_RATIO
: If--init_resource_reduction
is not specified,init_resource_reduction
=init_resource_reduction_ratio
*(pretrained model resource) (default: 0.025). -
-rd RESOURCE_REDUCTION_DECAY, --resource_reduction_decay RESOURCE_REDUCTION_DECAY
: For each iteration, target resource = current resource -init_resource_reduction
*(resource_reduction_decay
**(iteration-1)) (default: 0.96). -
-st SHORT_TERM_FINE_TUNE_ITERATION, --short_term_fine_tune_iteration SHORT_TERM_FINE_TUNE_ITERATION
: Short-term fine-tune iteration (default: 10). -
-lt LOOKUP_TABLE_PATH, --lookup_table_path LOOKUP_TABLE_PATH
: Path to lookup table. -
-dp DATASET_PATH, --dataset_path DATASET_PATH
: Path to dataset. -
-a ARCH, --arch ARCH network_utils
: Defines how networks are pruned, fine-tuned, and evaluated. If you want to use your own method, please see Customization and specify here. (default: alexnet) -
-si SAVE_INTERVAL, --save_interval SAVE_INTERVAL
: Interval of iterations that all pruned models at the same iteration will be saved. Use-1
to save only the best model at each iteration. Use1
to save all models at each iteration. (default: -1).
We provide a simple example of applying NetAdapt to a very small network:
sh scripts/netadapt_helloworld.sh
Detailed examples of applying NetAdapt to AlexNet/MobileNet on CIFAR-10 are shown here (AlexNet) and here (MobileNet).
If you want to apply the algorithm to different networks or even different tasks, please see the following Customization section.
To apply NetAdapt to differenct networks or different tasks, please follow the instructions:
-
Create your own
network_utils
python file (saidnetwork_utils_yourNetworkOrTask.py
) and place it undernetwork_utils
. -
Implement functions described in
network_utils_abstract.py
. -
As we provide an example of applying NetAdapt to AlexNet, you can also build your
network_utils
based onnetwork_utils_alexnet.py
:cd network_utils cp network_utils_alexnet.py ./network_utils_yourNetworkOrTask.py
-
Add
from .network_utils_yourNetworkOrTask import *
to__init__.py
, which is under the same directory. -
Modify
class networkUtils_alexnet(...)
in line 44 innetwork_utils_yourNetworkOrTask.py
toclass networkUtils_yourNetworkOrTask(...)
. -
Modify
def alexnet(...)
in line 325-326 to:def yourNetworkOrTask(model, input_data_shape, dataset_path, finetune_lr=1e-3): return networkUtils_yourNetworkOrTask(model, input_data_shape, dataset_path, finetune_lr)
-
Specify training/validation data loader, loss functions, optimizers, network architecture, training method, and evaluation method in
network_utils_yourNetworkOrTask.py
if there is any difference from the AlexNet example:-
Modify data loader and loss functionsin function
def __init__(...):
in line 52. -
Specify additive skip connections if there is any and modify
def simplify_network_def_based_on_constraint(...)
innetwork_utils_yourNetworkOrTask.py
. You can see how our implementation uses additive skip connections here. -
Modify training method (short-term finetune) in function
def fine_tune(...):
in line 245. -
Modify evaluation method in function
def evaluate(...):
in line 291.
You can see how these methods are utilized by the framework here.
-
-
Our current code base supports pruning
Conv2d
,ConvTranspose2d
, andLinear
with additive skip connection. If your network architecture is not supported, please modify this. If you want to use other metrics (resource type) to prune networks, please modify this. -
We can apply NetAdapt to different networks or tasks by using
--arch yourNetworkOrTask
inscripts/netadapt_alexnet-0.5mac.sh
. As for the values of other arguments, please see Usage. Generally, if you want to apply NetAdapt to a different task, you might changeinput_data_shape
. If your network architecture is very different from that of MobileNet, you would have to modify the values of--init_resource_reduction_ratio
and--resource_reduction_decay
to get a different resource reduction schedule.
If you use our code or method in your work, please consider citing the following:
@InProceedings{eccv_2018_yang_netadapt,
author = {Yang, Tien-Ju and Howard, Andrew and Chen, Bo and Zhang, Xiao and Go, Alec and Sandler, Mark and Sze, Vivienne and Adam, Hartwig},
title = {NetAdapt: Platform-Aware Neural Network Adaptation for Mobile Applications},
booktitle = {The European Conference on Computer Vision (ECCV)},
month = {September},
year = {2018}
}
Please direct any questions to the authors: Tien-Ju Yang ([email protected]) and Yi-Lun Liao ([email protected]).