Skip to content
Snippets Groups Projects
Commit b9f1c107 authored by Tien-Ju Yang's avatar Tien-Ju Yang
Browse files

Initial commit.

parents
No related branches found
No related tags found
No related merge requests found
LICENSE 0 → 100644
MIT License
Copyright (c) 2019 denru01
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md 0 → 100644
NetAdapt: Platform-Aware Neural Network Adaptation for Mobile Applications
============================
This repo contains the official Pytorch reimplementation of the paper "NetAdapt: Platform-Aware Neural Network Adaptation for Mobile Applications"
[[paper](http://openaccess.thecvf.com/content_ECCV_2018/papers/Tien-Ju_Yang_NetAdapt_Platform-Aware_Neural_ECCV_2018_paper.pdf)]
[[project](https://netadapt.mit.edu)]. The results in the paper were generated by the Tensorflow implementation from Google AI.
<p align="center">
<img src="fig/netadapt_fig.png" alt="photo not available" width="80%" height="80%">
</p>
## Summary ##
0. [Requirements](#requirements)
0. [Usage](#usage)
0. [Example](#example)
0. [Customization](#customization)
0. [Citation](#citation)
<!---
## Docs ##
`constants.py` : defines constants and some layers that can be simplified by this algorithm.
`worker.py` : called by `master.py` to simplify a certain block of the network based on resource constraints and then run short-term finetune of simplified models. Several workers may execute at the same time.
`network_utils_abstract.py` : defines abstract base class for networks.
`network_utils_depth_mobilenetv1` : an example defining mobilenet, its correspoding loss function, data loaders.
`functions.py` : defines some utility functions like computing MACs, getting feature map sizes that would be used `network_utils`.
--->
## Requirements ##
The code base is tested with the following setting:
1. Python 3.7.0
2. CUDA 10.0
3. Pytorch 1.2.0
4. torchvision 0.4.0
5. numpy 1.17.0
6. scipy 1.3.1
First clone the repo in the directory you want to work:
git clone ...
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.
## Usage ##
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 in `working_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 specify `working_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 (under `working_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 support `FLOPS`, `WEIGHTS`, and `LATENCY` (device `cuda:0`). If you want to add other resource
types, please modify `def compute_resource(...)` in `network_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**](#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.
Use `1` to save all models at each iteration. (default: -1).
## Example ##
<!---
We provide an example of applying **NetAdapt** to **AlexNet** on **CIFAR-10**.
If you want to apply the algorithm to different networks or even different tasks,
please see the following [**Customization**](#customization) section.
The example consists of the following four steps: 1. training AlexNet on CIFAR-10; 2. measuring latency; 3. applying NetAdapt; 4. evaluation using adapted models.
1. **Training AlexNet on CIFAR-10.**
Train from scratch:
```bash
python train.py data/ --dir models/alexnet/no-pretrain/checkpoint.pth.tar --arch alexnet
```
Use Imagenet-pretrained model:
```bash
python train.py data/ --dir models/alexnet/pretrain/checkpoint.pth.tar --pretrained --lr 0.01 --arch alexnet
```
Evaluation:
```bash
python eval.py data/ --dir models/alexnet/no-pretrain/checkpoint.pth.tar --arch alexnet
```
```bash
python train.py data/ --dir models/mobilenet/model.pth.tar --arch mobilenet
```
2. **Measuring Latency**
Here we build the latency lookup table for `cuda:0` device:
```bash
python build_lookup_table.py --dir latency_lut/lut_alexnet.pkl --arch alexnet
```
It measures latency of different layers contained in the network (i.e. **AlexNet** here).
For conv layers, the sampled numbers of feature channels are multiples of `MIN_CONV_FEATURE_SIZE`.
For fc layers, the sampled numbers of features are multiples of `MIN_FC_FEATURE_SIZE`.
3. **Applying NetAdapt**
Modify which GPUs will be utilized (`-gp`) in `netadapt_alexnet.sh` and run the script to apply NetAdapt to a pretrained model:
```bash
sh netadapt_alexnet.sh
```
After obtaining the adapted model, we need to finetune the model:
```bash
python train.py data/ --arch alexnet --resume models/alexnet/no-pretrain/prune-by-mac/iter_100_best_model.pth.tar --dir models/alexnet/no-pretrain/checkpoint-adapt.pth.tar
```
<p align="center">
<img src="fig/netadapt_algo.png" alt="photo not available" width="90%" height="90%">
</p>
4. **Evaluation Using Adapted Models**
After applying NetAdapt to a pretrained model, we can evaluate this adapted model using:
```bash
python eval.py data/ --dir models/alexnet/no-pretrain/checkpoint-adapt.pth.tar --arch alexnet
```
The adapted model can be restored **without modifying the orignal python file**.
--->
We provide a simple example of applying **NetAdapt** to a very small [network](nets/helloworld.py):
sh scripts/netadapt_helloworld.sh
Detailed examples of applying **NetAdapt** to **AlexNet**/**MobileNet** on **CIFAR-10** are shown [**here (AlexNet)**](docs/alexnet/README.md) and [**here (MobileNet)**](docs/mobilenet/README.md).
<p align="center">
<img src="fig/netadapt_algo.png" alt="photo not available" width="90%" height="90%">
</p>
If you want to apply the algorithm to different networks or even different tasks,
please see the following [**Customization**](#customization) section.
## Customization ##
To apply NetAdapt to differenct networks or different tasks, please follow the instructions:
1. Create your own `network_utils` python file (said `network_utils_yourNetworkOrTask.py`) and place it under `network_utils`.
2. Implement functions described in [`network_utils_abstract.py`](network_utils/network_utils_abstract.py).
3. As we provide an example of applying NetAdapt to AlexNet, you can also build your `network_utils` based on [`network_utils_alexnet.py`](network_utils/network_utils_alexnet.py):
```bash
cd network_utils
cp network_utils_alexnet.py ./network_utils_yourNetworkOrTask.py
```
4. Add `from .network_utils_yourNetworkOrTask import *` to `__init__.py`, which is under [the same directory](network_utils/__init__.py).
5. Modify `class networkUtils_alexnet(...)` in [line 44](network_utils/network_utils_alexnet.py#L44) in `network_utils_yourNetworkOrTask.py` to `class networkUtils_yourNetworkOrTask(...)`.
6. Modify `def alexnet(...)` in [line 325-326](network_utils/network_utils_alexnet.py#L325-L326) to:
```bash
def yourNetworkOrTask(model, input_data_shape, dataset_path, finetune_lr=1e-3):
return networkUtils_yourNetworkOrTask(model, input_data_shape, dataset_path, finetune_lr)
```
7. 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](network_utils/network_utils_alexnet.py#L52-L125).
- Specify additive skip connections if there is any and modify `def simplify_network_def_based_on_constraint(...)` in `network_utils_yourNetworkOrTask.py`.
You can see how our implementation uses additive skip connections [here](functions.py#L543-L549).
- Modify training method (short-term finetune) in function `def fine_tune(...):` in [line 245](network_utils/network_utils_alexnet.py#L245-L288).
- Modify evaluation method in function `def evaluate(...):` in [line 291](network_utils/network_utils_alexnet.py#L291-L322).
You can see how these methods are utilized by the framework [here](worker.py#L39-L62).
8. Our current code base supports pruning `Conv2d`, `ConvTranspose2d`, and `Linear` with additive skip connection.
If your network architecture is not supported, please modify [this](network_utils/network_utils_alexnet.py#L142-L199).
If you want to use other metrics (resource type) to prune networks, please modify [this](network_utils/network_utils_alexnet.py#L234-L238).
9. We can apply NetAdapt to different networks or tasks by using `--arch yourNetworkOrTask` in `scripts/netadapt_alexnet-0.5mac.sh`.
As for the values of other arguments, please see [**Usage**](#usage).
Generally, if you want to apply NetAdapt to a different task, you might change `input_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.
## Citation
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 (tjy@mit.edu) and Yi-Lun Liao (ylliao@mit.edu).
import torch
import torch.nn as nn
import nets as models
import functions as fns
from argparse import ArgumentParser
model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__")
and callable(models.__dict__[name]))
NUM_CLASSES = 10
INPUT_DATA_SHAPE = (3, 224, 224)
'''
`MIN_CONV_FEATURE_SIZE`: The sampled size of feature maps of layers (conv layer)
along channel dimmension are multiples of 'MIN_CONV_FEATURE_SIZE'.
`MIN_FC_FEATURE_SIZE`: The sampled size of features of FC layers are
multiples of 'MIN_FC_FEATURE_SIZE'.
'''
MIN_CONV_FEATURE_SIZE = 8
MIN_FC_FEATRE_SIZE = 64
'''
`MEASURE_LATENCY_BATCH_SIZE`: the batch size of input data
when running forward functions to measure latency.
`MEASURE_LATENCY_SAMPLE_TIMES`: the number of times to run the forward function of
a layer in order to get its latency.
'''
MEASURE_LATENCY_BATCH_SIZE = 128
MEASURE_LATENCY_SAMPLE_TIMES = 500
arg_parser = ArgumentParser(description='Build latency lookup table')
arg_parser.add_argument('--dir', metavar='DIR', default='latency_lut/lut_alexnet.pkl',
help='path to saving lookup table')
arg_parser.add_argument('-a', '--arch', metavar='ARCH', default='alexnet',
choices=model_names,
help='model architecture: ' +
' | '.join(model_names) +
' (default: alexnet)')
if __name__ == '__main__':
args = arg_parser.parse_args()
print(args)
build_lookup_table = True
lookup_table_path = args.dir
model_arch = args.arch
print('Load', model_arch)
print('--------------------------------------')
model = models.__dict__[model_arch](num_classes=NUM_CLASSES)
network_def = fns.get_network_def_from_model(model, INPUT_DATA_SHAPE)
for layer_name, layer_properties in network_def.items():
print(layer_name)
print(' ', layer_properties, '\n')
print('-------------------------------------------')
num_w = fns.compute_resource(network_def, 'WEIGHTS')
flops = fns.compute_resource(network_def, 'FLOPS')
num_param = fns.compute_resource(network_def, 'WEIGHTS')
print('Number of FLOPs: ', flops)
print('Number of weights: ', num_w)
print('Number of parameters: ', num_param)
print('-------------------------------------------')
model = model.cuda()
print('Building latency lookup table for',
torch.cuda.get_device_name())
if build_lookup_table:
fns.build_latency_lookup_table(network_def, lookup_table_path=lookup_table_path,
min_fc_feature_size=MIN_FC_FEATRE_SIZE,
min_conv_feature_size=MIN_CONV_FEATURE_SIZE,
measure_latency_batch_size=MEASURE_LATENCY_BATCH_SIZE,
measure_latency_sample_times=MEASURE_LATENCY_SAMPLE_TIMES,
verbose=True)
print('-------------------------------------------')
print('Finish building latency lookup table.')
print(' Device:', torch.cuda.get_device_name())
print(' Model: ', model_arch)
print('-------------------------------------------')
latency = fns.compute_resource(network_def, 'LATENCY', lookup_table_path)
print('Computed latency: ', latency)
latency = fns.measure_latency(model,
[MEASURE_LATENCY_BATCH_SIZE, *INPUT_DATA_SHAPE])
print('Exact latency: ', latency)
\ No newline at end of file
# Master-related filenames.
MASTER_MODEL_FILENAME_TEMPLATE = 'iter_{}_best_model.pth.tar'
# Worker-related filenames.
WORKER_MODEL_FILENAME_TEMPLATE = 'iter_{}_block_{}_model.pth.tar'
WORKER_ACCURACY_FILENAME_TEMPLATE = 'iter_{}_block_{}_accuracy.txt'
WORKER_RESOURCE_FILENAME_TEMPLATE = 'iter_{}_block_{}_resource.txt'
WORKER_LOG_FILENAME_TEMPLATE = 'iter_{}_block_{}_log.txt'
WORKER_FINISH_FILENAME_TEMPLATE = 'iter_{}_block_{}_finish.signal'
\ No newline at end of file
# Define constants.
STRING_SEPARATOR = '.'
# Define layer types.
CONV_LAYER_TYPES = ['Conv2d', 'ConvTranspose2d']
FC_LAYER_TYPES = ['Linear']
BNORM_LAYER_TYPES = ['BatchNorm2d']
# Define data types.
WEIGHTSTRING = 'weight'
BIASSTRING = 'bias'
RUNNING_MEANSTRING = 'running_mean'
RUNNING_VARSTRING = 'running_var'
NUM_BATCHES_TRACKED = 'num_batches_tracked'
# Define keys.
KEY_LAYER_TYPE_STR = 'layer_type_str' #(e.g. Linear, Conv2d)
KEY_IS_DEPTHWISE = 'is_depthwise'
KEY_NUM_IN_CHANNELS = 'num_in_channels'
KEY_NUM_OUT_CHANNELS = 'num_out_channels'
KEY_KERNEL_SIZE = 'kernel_size'
KEY_STRIDE = 'stride'
KEY_PADDING = 'padding'
KEY_GROUPS = 'groups'
KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR = 'before_squared_pixel_shuffle_factor'
KEY_AFTER_SQUSRED_PIXEL_SHUFFLE_FACTOR = 'after_squared_pixel_shuffle_factor'
KEY_INPUT_FEATURE_MAP_SIZE = 'input_feature_map_size'
KEY_OUTPUT_FEATURE_MAP_SIZE = 'output_feature_map_size'
KEY_MODEL = 'model'
KEY_LATENCY = 'latency'
\ No newline at end of file
1. **Training AlexNet on CIFAR-10.**
Training:
```bash
python train.py data/ --dir models/alexnet/model.pth.tar --arch alexnet
```
Evaluation:
```bash
python eval.py data/ --dir models/alexnet/model.pth.tar --arch alexnet
```
One trained model can be found [here](https://drive.google.com/file/d/1GwugqlSl5ogRvJ2F2-4uz1fXBzt36ahB/view?usp=sharing).
2. **Measuring Latency**
Here we build the latency lookup table for `cuda:0` device:
```bash
python build_lookup_table.py --dir latency_lut/lut_alexnet.pkl --arch alexnet
```
It measures latency of different layers contained in the network (i.e. **AlexNet** here).
For conv layers, the sampled numbers of feature channels are multiples of `MIN_CONV_FEATURE_SIZE`.
For fc layers, the sampled numbers of features are multiples of `MIN_FC_FEATURE_SIZE`.
3. **Applying NetAdapt**
Modify which GPUs will be utilized (`-gp`) in `netadapt_alexnet-0.5mac.sh` and run the script to apply NetAdapt to a pretrained model:
```bash
sh scripts/netadapt_alexnet-0.5mac.sh
```
You can see how the model is simplified at each iteration in `models/alexnet/prune-by-mac/master/history.txt` and
select the one that satisfies the constraints to run long-term fine-tune.
After obtaining the adapted model, we need to finetune the model (here we select the one after 18 iterations):
```bash
python train.py data/ --arch alexnet --resume models/alexnet/prune-by-mac/master/iter_18_best_model.pth.tar --dir models/alexnet/prune-by-mac/master/finetune_model.pth.tar --lr 0.001
```
<p align="center">
<img src="../../fig/netadapt_algo.png" alt="photo not available" width="90%" height="90%">
</p>
If you want to get a model with 50% latency, please run:
```bash
sh scripts/netadapt_alexnet-0.5latency.sh
```
4. **Evaluation Using Adapted Models**
After applying NetAdapt to a pretrained model, we can evaluate this adapted model using:
```bash
python eval.py data/ --dir models/alexnet/prune-by-mac/master/finetune_model.pth.tar --arch alexnet
```
The adapted model can be restored **without modifying the orignal python file**.
We provide one adapted model [here](https://drive.google.com/file/d/1VH9c2orF2W0P21gD8NrdTdvwP_uJJgYA/view?usp=sharing).
\ No newline at end of file
1. **Training MobileNet on CIFAR-10.**
Training:
```bash
python train.py data/ --dir models/mobilenet/model.pth.tar --arch mobilenet
```
Evaluation:
```bash
python eval.py data/ --dir models/mobilenet/model.pth.tar --arch mobilenet
```
One trained model can be found [here](https://drive.google.com/file/d/1jtRZOHK1daRTKD4jYu4US84lf8YjqEdJ/view?usp=sharing).
2. **Measuring Latency**
Here we build the latency lookup table for `cuda:0` device:
```bash
python build_lookup_table.py --dir latency_lut/lut_mobilenet.pkl --arch mobilenet
```
It measures latency of different layers contained in the network (i.e. **MobileNet** here).
For conv layers, the sampled numbers of feature channels are multiples of `MIN_CONV_FEATURE_SIZE`.
For fc layers, the sampled numbers of features are multiples of `MIN_FC_FEATURE_SIZE`.
3. **Applying NetAdapt**
Modify which GPUs will be utilized (`-gp`) in `netadapt_mobilenet-0.5mac.sh` and run the script to apply NetAdapt to a pretrained model:
```bash
sh scripts/netadapt_mobilenet-0.5mac.sh
```
You can see how the model is simplified at each iteration in `models/mobilenet/prune-by-mac/master/history.txt` and
select the one that satisfies the constraints to run long-term fine-tune.
After obtaining the adapted model, we need to finetune the model (here we select the one after 28 iterations):
```bash
python train.py data/ --arch mobilenet --resume models/mobilenet/prune-by-mac/master/iter_28_best_model.pth.tar --dir models/mobilenet/prune-by-mac/master/finetune_model.pth.tar --lr 0.001
```
<p align="center">
<img src="../../fig/netadapt_algo.png" alt="photo not available" width="90%" height="90%">
</p>
If you want to get a model with 50% latency, please run:
```bash
sh scripts/netadapt_mobilenet-0.5latency.sh
```
4. **Evaluation Using Adapted Models**
After applying NetAdapt to a pretrained model, we can evaluate this adapted model using:
```bash
python eval.py data/ --dir models/mobilenet/prune-by-mac/master/finetune_model.pth.tar --arch mobilenet
```
The adapted model can be restored **without modifying the orignal python file**.
We provide one adapted model [here](https://drive.google.com/file/d/1wkQPolgv34ESb0gyeYdygbuNuNYhIqyx/view?usp=sharing).
\ No newline at end of file
eval.py 0 → 100644
from argparse import ArgumentParser
import os
import time
import math
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.backends.cudnn as cudnn
import pickle
import nets as models
import functions as fns
_NUM_CLASSES = 10
model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__")
and callable(models.__dict__[name]))
def compute_topk_accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def compute_accuracy(output, target):
output = output.argmax(dim=1)
acc = 0.0
acc = torch.sum(target == output).item()
acc = acc/output.size(0)*100
return acc
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def get_avg(self):
return self.avg
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def eval(test_loader, model, args):
batch_time = AverageMeter()
acc = AverageMeter()
# switch to eval mode
model.eval()
end = time.time()
for i, (images, target) in enumerate(test_loader):
if not args.no_cuda:
images = images.cuda()
target = target.cuda()
output = model(images)
batch_acc = compute_accuracy(output, target)
acc.update(batch_acc, images.size(0))
batch_time.update(time.time() - end)
end = time.time()
# Update statistics
estimated_time_remained = batch_time.get_avg()*(len(test_loader)-i-1)
fns.update_progress(i, len(test_loader),
ESA='{:8.2f}'.format(estimated_time_remained)+'s',
acc='{:4.2f}'.format(float(batch_acc))
)
print()
print('Test accuracy: {:4.2f}% (time = {:8.2f}s)'.format(
float(acc.get_avg()), batch_time.get_avg()*len(test_loader)))
print('===================================================================')
return float(acc.get_avg())
if __name__ == '__main__':
# Parse the input arguments.
arg_parser = ArgumentParser()
arg_parser.add_argument('data', metavar='DIR', help='path to dataset')
arg_parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
help='number of data loading workers (default: 4)')
arg_parser.add_argument('-a', '--arch', metavar='ARCH', default='alexnet',
choices=model_names,
help='model architecture: ' +
' | '.join(model_names) +
' (default: alexnet)')
arg_parser.add_argument('-b', '--batch-size', default=128, type=int,
metavar='N',
help='batch size (default: 128)')
arg_parser.add_argument('--dir', type=str, default='models/', dest='save_dir',
help='path to save models (default: models/')
arg_parser.add_argument('--no-cuda', action='store_true', default=False, dest='no_cuda',
help='disables training on GPU')
args = arg_parser.parse_args()
print(args)
# Data loader
test_dataset = datasets.CIFAR10(root=args.data, train=False, download=True,
transform=transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
]))
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# Network
model_arch = args.arch
cudnn.benchmark = True
num_classes = _NUM_CLASSES
model = models.__dict__[model_arch](num_classes=num_classes)
if not args.no_cuda:
model = model.cuda()
# Evaluation
filename = os.path.join(args.save_dir)
model = torch.load(filename)
print(model)
best_acc = eval(test_loader, model, args)
print('Testing accuracy:', best_acc)
\ No newline at end of file
fig/netadapt_algo.png

171 KiB

fig/netadapt_fig.png

97.8 KiB

This diff is collapsed.
This diff is collapsed.
File added
File added
from .alexnet import *
from .mobilenet import *
from .helloworld import *
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment