Skip to content

GitLab

Commit 1011377c authored by qianyj's avatar qianyj
Browse files

the source code of NNI for DCU

parent abc22158
Hide whitespace changes
Inline Side-by-side
Showing with 3098 additions and 0 deletions
examples/model_compress/pruning/comparison_of_pruners/cifar10/comparison_result_vgg16.json 0 → 100644
View file @ 1011377c
{
"L1FilterPruner": [
{
"sparsity": 0.1,
"params": 12187336.0,
"flops": 256252606.0,
"performance": 0.9344
},
{
"sparsity": 0.2,
"params": 9660216.0,
"flops": 203049930.0,
"performance": 0.9371
},
{
"sparsity": 0.3,
"params": 7435417.0,
"flops": 155477470.0,
"performance": 0.9341
},
{
"sparsity": 0.4,
"params": 5493954.0,
"flops": 114721578.0,
"performance": 0.9317
},
{
"sparsity": 0.5,
"params": 3820010.0,
"flops": 79155722.0,
"performance": 0.9309
},
{
"sparsity": 0.6,
"params": 2478632.0,
"flops": 51618494.0,
"performance": 0.9229
},
{
"sparsity": 0.7,
"params": 1420600.0,
"flops": 29455306.0,
"performance": 0.9031
},
{
"sparsity": 0.8,
"params": 658553.0,
"flops": 13290974.0,
"performance": 0.8756
},
{
"sparsity": 0.9,
"params": 186178.0,
"flops": 3574570.0,
"performance": 0.8145
},
{
"sparsity": 0.95,
"params": 58680.0,
"flops": 1050570.0,
"performance": 0.6983
},
{
"sparsity": 0.975,
"params": 23408.0,
"flops": 329918.0,
"performance": 0.5573
}
],
"L2FilterPruner": [
{
"sparsity": 0.1,
"params": 12187336.0,
"flops": 256252606.0,
"performance": 0.9357
},
{
"sparsity": 0.2,
"params": 9660216.0,
"flops": 203049930.0,
"performance": 0.9355
},
{
"sparsity": 0.3,
"params": 7435417.0,
"flops": 155477470.0,
"performance": 0.9337
},
{
"sparsity": 0.4,
"params": 5493954.0,
"flops": 114721578.0,
"performance": 0.9308
},
{
"sparsity": 0.5,
"params": 3820010.0,
"flops": 79155722.0,
"performance": 0.9285
},
{
"sparsity": 0.6,
"params": 2478632.0,
"flops": 51618494.0,
"performance": 0.9208
},
{
"sparsity": 0.7,
"params": 1420600.0,
"flops": 29455306.0,
"performance": 0.909
},
{
"sparsity": 0.8,
"params": 658553.0,
"flops": 13290974.0,
"performance": 0.8698
},
{
"sparsity": 0.9,
"params": 186178.0,
"flops": 3574570.0,
"performance": 0.8203
},
{
"sparsity": 0.95,
"params": 58680.0,
"flops": 1050570.0,
"performance": 0.7063
},
{
"sparsity": 0.975,
"params": 23408.0,
"flops": 329918.0,
"performance": 0.5455
}
],
"FPGMPruner": [
{
"sparsity": 0.1,
"params": 12187336.0,
"flops": 256252606.0,
"performance": 0.937
},
{
"sparsity": 0.2,
"params": 9660216.0,
"flops": 203049930.0,
"performance": 0.936
},
{
"sparsity": 0.3,
"params": 7435417.0,
"flops": 155477470.0,
"performance": 0.9359
},
{
"sparsity": 0.4,
"params": 5493954.0,
"flops": 114721578.0,
"performance": 0.9302
},
{
"sparsity": 0.5,
"params": 3820010.0,
"flops": 79155722.0,
"performance": 0.9233
},
{
"sparsity": 0.6,
"params": 2478632.0,
"flops": 51618494.0,
"performance": 0.922
},
{
"sparsity": 0.7,
"params": 1420600.0,
"flops": 29455306.0,
"performance": 0.9022
},
{
"sparsity": 0.8,
"params": 658553.0,
"flops": 13290974.0,
"performance": 0.8794
},
{
"sparsity": 0.9,
"params": 186178.0,
"flops": 3574570.0,
"performance": 0.8276
},
{
"sparsity": 0.95,
"params": 58680.0,
"flops": 1050570.0,
"performance": 0.6967
},
{
"sparsity": 0.975,
"params": 23408.0,
"flops": 329918.0,
"performance": 0.3683
}
],
"NetAdaptPruner": [
{
"sparsity": 0.1,
"params": 13492098.0,
"flops": 308484330.0,
"performance": 0.9376
},
{
"sparsity": 0.2,
"params": 11998408.0,
"flops": 297641410.0,
"performance": 0.9374
},
{
"sparsity": 0.3,
"params": 10504344.0,
"flops": 281928834.0,
"performance": 0.9369
},
{
"sparsity": 0.4,
"params": 8263221.0,
"flops": 272964342.0,
"performance": 0.9382
},
{
"sparsity": 0.5,
"params": 6769885.0,
"flops": 249070966.0,
"performance": 0.9388
},
{
"sparsity": 0.6,
"params": 6022137.0,
"flops": 237106998.0,
"performance": 0.9383
},
{
"sparsity": 0.7,
"params": 4526754.0,
"flops": 222152490.0,
"performance": 0.936
},
{
"sparsity": 0.8,
"params": 3032759.0,
"flops": 162401210.0,
"performance": 0.9362
}
],
"SimulatedAnnealingPruner": [
{
"sparsity": 0.1,
"params": 12691704.0,
"flops": 301467870.0,
"performance": 0.9366
},
{
"sparsity": 0.2,
"params": 10318461.0,
"flops": 275724450.0,
"performance": 0.9362
},
{
"sparsity": 0.3,
"params": 8217127.0,
"flops": 246321046.0,
"performance": 0.9371
},
{
"sparsity": 0.4,
"params": 6458368.0,
"flops": 232948294.0,
"performance": 0.9378
},
{
"sparsity": 0.5,
"params": 4973079.0,
"flops": 217675254.0,
"performance": 0.9362
},
{
"sparsity": 0.6,
"params": 3131526.0,
"flops": 151576878.0,
"performance": 0.9347
},
{
"sparsity": 0.7,
"params": 1891036.0,
"flops": 76575574.0,
"performance": 0.9289
},
{
"sparsity": 0.8,
"params": 1170751.0,
"flops": 107532322.0,
"performance": 0.9325
},
{
"sparsity": 0.9,
"params": 365978.0,
"flops": 46241354.0,
"performance": 0.9167
},
{
"sparsity": 0.95,
"params": 167089.0,
"flops": 38589922.0,
"performance": 0.7746
},
{
"sparsity": 0.975,
"params": 96779.0,
"flops": 26838230.0,
"performance": 0.1
}
],
"AutoCompressPruner": [
{
"sparsity": 0.1,
"params": 12460277.0,
"flops": 290311730.0,
"performance": 0.9352
},
{
"sparsity": 0.2,
"params": 10138147.0,
"flops": 269180938.0,
"performance": 0.9324
},
{
"sparsity": 0.3,
"params": 8033350.0,
"flops": 241789714.0,
"performance": 0.9357
},
{
"sparsity": 0.4,
"params": 6105156.0,
"flops": 213573294.0,
"performance": 0.9367
},
{
"sparsity": 0.5,
"params": 4372604.0,
"flops": 185826362.0,
"performance": 0.9387
},
{
"sparsity": 0.6,
"params": 3029629.0,
"flops": 166285498.0,
"performance": 0.9334
},
{
"sparsity": 0.7,
"params": 1897060.0,
"flops": 134897806.0,
"performance": 0.9359
},
{
"sparsity": 0.8,
"params": 1145509.0,
"flops": 111766450.0,
"performance": 0.9334
},
{
"sparsity": 0.9,
"params": 362546.0,
"flops": 50777246.0,
"performance": 0.9261
},
{
"sparsity": 0.95,
"params": 149735.0,
"flops": 39201770.0,
"performance": 0.8924
},
{
"sparsity": 0.975,
"params": 45378.0,
"flops": 13213974.0,
"performance": 0.8193
}
]
}
\ No newline at end of file
examples/model_compress/pruning/config.yml 0 → 100644
View file @ 1011377c
searchSpace:
sparsity:
_type: choice
_value: [0.25, 0.5, 0.75]
pruner:
_type: choice
_value: ['slim', 'l2filter', 'fpgm', 'apoz']
model:
_type: choice
_value: ['vgg16', 'vgg19']
trainingService:
platform: local
trialCodeDirectory: .
trialCommand: python3 basic_pruners_torch.py --nni
trialConcurrency: 1
trialGpuNumber: 0
tuner:
name: GridSearch
examples/model_compress/pruning/finetune_kd_torch.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
'''
NNI exmaple for fine-tuning the pruned model with KD.
Run basic_pruners_torch.py first to get the masks of the pruned model. Then pass the mask as argument for model speedup. The compressed model is further used for fine-tuning.
'''
import argparse
import os
import sys
from copy import deepcopy
import torch
import torch.nn as nn
import torch.nn.functional as F
from nni.compression.pytorch import ModelSpeedup
from torch.optim.lr_scheduler import MultiStepLR
from basic_pruners_torch import get_data
from pathlib import Path
sys.path.append(str(Path(__file__).absolute().parents[1] / 'models'))
from mnist.lenet import LeNet
from cifar10.vgg import VGG
class DistillKL(nn.Module):
"""Distilling the Knowledge in a Neural Network"""
def __init__(self, T):
super(DistillKL, self).__init__()
self.T = T
def forward(self, y_s, y_t):
p_s = F.log_softmax(y_s/self.T, dim=1)
p_t = F.softmax(y_t/self.T, dim=1)
loss = F.kl_div(p_s, p_t, size_average=False) * (self.T**2) / y_s.shape[0]
return loss
def get_dummy_input(args, device):
if args.dataset == 'mnist':
dummy_input = torch.randn([args.test_batch_size, 1, 28, 28]).to(device)
elif args.dataset in ['cifar10', 'imagenet']:
dummy_input = torch.randn([args.test_batch_size, 3, 32, 32]).to(device)
return dummy_input
def get_model_optimizer_scheduler(args, device, test_loader, criterion):
if args.model == 'LeNet':
model = LeNet().to(device)
elif args.model == 'vgg16':
model = VGG(depth=16).to(device)
elif args.model == 'vgg19':
model = VGG(depth=19).to(device)
else:
raise ValueError("model not recognized")
# In this example, we set the architecture of teacher and student to be the same. It is feasible to set a different teacher architecture.
if args.teacher_model_dir is None:
raise NotImplementedError('please load pretrained teacher model first')
else:
model.load_state_dict(torch.load(args.teacher_model_dir))
best_acc = test(args, model, device, criterion, test_loader)
model_t = deepcopy(model)
model_s = deepcopy(model)
if args.student_model_dir is not None:
# load the pruned student model checkpoint
model_s.load_state_dict(torch.load(args.student_model_dir))
dummy_input = get_dummy_input(args, device)
m_speedup = ModelSpeedup(model_s, dummy_input, args.mask_path, device)
m_speedup.speedup_model()
module_list = nn.ModuleList([])
module_list.append(model_s)
module_list.append(model_t)
# setup opotimizer for fine-tuning studeng model
optimizer = torch.optim.SGD(model_s.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
scheduler = MultiStepLR(
optimizer, milestones=[int(args.fine_tune_epochs*0.5), int(args.fine_tune_epochs*0.75)], gamma=0.1)
print('Pretrained teacher model acc:', best_acc)
return module_list, optimizer, scheduler
def train(args, models, device, train_loader, criterion, optimizer, epoch):
# model.train()
model_s = models[0].train()
model_t = models[-1].eval()
cri_cls = criterion
cri_kd = DistillKL(args.kd_T)
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output_s = model_s(data)
output_t = model_t(data)
loss_cls = cri_cls(output_s, target)
loss_kd = cri_kd(output_s, output_t)
loss = loss_cls + loss_kd
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
if args.dry_run:
break
def test(args, model, device, criterion, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
acc = 100 * correct / len(test_loader.dataset)
print('Test Loss: {} Accuracy: {}%\n'.format(
test_loss, acc))
return acc
def main(args):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs(args.experiment_data_dir, exist_ok=True)
# prepare model and data
train_loader, test_loader, criterion = get_data(args.dataset, args.data_dir, args.batch_size, args.test_batch_size)
models, optimizer, scheduler = get_model_optimizer_scheduler(args, device, test_loader, criterion)
best_top1 = 0
if args.test_only:
test(args, models[0], device, criterion, test_loader)
print('start fine-tuning...')
for epoch in range(args.fine_tune_epochs):
print('# Epoch {} #'.format(epoch))
train(args, models, device, train_loader, criterion, optimizer, epoch)
scheduler.step()
# test student only
top1 = test(args, models[0], device, criterion, test_loader)
if top1 > best_top1:
best_top1 = top1
torch.save(models[0].state_dict(), os.path.join(args.experiment_data_dir, 'model_trained.pth'))
print('Model trained saved to %s' % args.experiment_data_dir)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='PyTorch Example for model comporession')
# dataset and model
parser.add_argument('--dataset', type=str, default='cifar10',
help='dataset to use, mnist, cifar10 or imagenet')
parser.add_argument('--data-dir', type=str, default='./data/',
help='dataset directory')
parser.add_argument('--model', type=str, default='vgg16',
choices=['LeNet', 'vgg16' ,'vgg19', 'resnet18'],
help='model to use')
parser.add_argument('--teacher-model-dir', type=str, default=None,
help='path to the pretrained teacher model checkpoint')
parser.add_argument('--mask-path', type=str, default=None,
help='path to the pruned student model mask file')
parser.add_argument('--student-model-dir', type=str, default=None,
help='path to the pruned student model checkpoint')
parser.add_argument('--batch-size', type=int, default=128,
help='input batch size for training')
parser.add_argument('--test-batch-size', type=int, default=200,
help='input batch size for testing')
parser.add_argument('--fine-tune-epochs', type=int, default=160,
help='epochs to fine tune')
parser.add_argument('--experiment-data-dir', type=str, default='./experiment_data',
help='For saving output checkpoints')
parser.add_argument('--log-interval', type=int, default=100, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--dry-run', action='store_true', default=False,
help='quickly check a single pass')
parser.add_argument('--test-only', action='store_true', default=False,
help='run test only')
# knowledge distillation
parser.add_argument('--kd_T', type=float, default=4,
help='temperature for KD distillation')
args = parser.parse_args()
main(args)
examples/model_compress/pruning/lottery_torch_mnist_fc.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
'''
NNI exmaple for reproducing Lottery Ticket Hypothesis.
'''
import argparse
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from nni.algorithms.compression.pytorch.pruning import LotteryTicketPruner
class fc1(nn.Module):
def __init__(self, num_classes=10):
super(fc1, self).__init__()
self.classifier = nn.Sequential(
nn.Linear(28 * 28, 300),
nn.ReLU(inplace=True),
nn.Linear(300, 100),
nn.ReLU(inplace=True),
nn.Linear(100, num_classes),
)
def forward(self, x):
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
def train(model, train_loader, optimizer, criterion):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.train()
for batch_idx, (imgs, targets) in enumerate(train_loader):
optimizer.zero_grad()
imgs, targets = imgs.to(device), targets.to(device)
output = model(imgs)
train_loss = criterion(output, targets)
train_loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print('{:2.0f}% Loss {}'.format(
100 * batch_idx / len(train_loader), train_loss.item()))
return train_loss.item()
def test(model, test_loader, criterion):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
return accuracy
if __name__ == '__main__':
"""
THE LOTTERY TICKET HYPOTHESIS: FINDING SPARSE, TRAINABLE NEURAL NETWORKS (https://arxiv.org/pdf/1803.03635.pdf)
The Lottery Ticket Hypothesis. A randomly-initialized, dense neural network contains a subnetwork that is
initialized such that—when trained in isolation—it can match the test accuracy of the original network after
training for at most the same number of iterations.
Identifying winning tickets. We identify a winning ticket by training a network and pruning its
smallest-magnitude weights. The remaining, unpruned connections constitute the architecture of the
winning ticket. Unique to our work, each unpruned connection’s value is then reset to its initialization
from original network before it was trained. This forms our central experiment:
1. Randomly initialize a neural network f(x; θ0) (where θ0 ∼ Dθ).
2. Train the network for j iterations, arriving at parameters θj .
3. Prune p% of the parameters in θj , creating a mask m.
4. Reset the remaining parameters to their values in θ0, creating the winning ticket f(x; m θ0).
As described, this pruning approach is one-shot: the network is trained once, p% of weights are
pruned, and the surviving weights are reset. However, in this paper, we focus on iterative pruning,
which repeatedly trains, prunes, and resets the network over n rounds; each round prunes p**(1/n) % of
the weights that survive the previous round. Our results show that iterative pruning finds winning tickets
that match the accuracy of the original network at smaller sizes than does one-shot pruning.
"""
parser = argparse.ArgumentParser()
parser.add_argument("--train_epochs", type=int, default=10, help="training epochs")
args = parser.parse_args()
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
traindataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
testdataset = datasets.MNIST('./data', train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(traindataset, batch_size=60, shuffle=True, num_workers=0, drop_last=False)
test_loader = torch.utils.data.DataLoader(testdataset, batch_size=60, shuffle=False, num_workers=0, drop_last=True)
model = fc1().to("cuda" if torch.cuda.is_available() else "cpu")
optimizer = torch.optim.Adam(model.parameters(), lr=1.2e-3)
criterion = nn.CrossEntropyLoss()
# Record the random intialized model weights
orig_state = copy.deepcopy(model.state_dict())
# train the model to get unpruned metrics
for epoch in range(args.train_epochs):
train(model, train_loader, optimizer, criterion)
orig_accuracy = test(model, test_loader, criterion)
print('unpruned model accuracy: {}'.format(orig_accuracy))
# reset model weights and optimizer for pruning
model.load_state_dict(orig_state)
optimizer = torch.optim.Adam(model.parameters(), lr=1.2e-3)
# Prune the model to find a winning ticket
configure_list = [{
'prune_iterations': 5,
'sparsity': 0.96,
'op_types': ['default']
}]
pruner = LotteryTicketPruner(model, configure_list, optimizer)
pruner.compress()
best_accuracy = 0.
best_state_dict = None
for i in pruner.get_prune_iterations():
pruner.prune_iteration_start()
loss = 0
accuracy = 0
for epoch in range(args.train_epochs):
loss = train(model, train_loader, optimizer, criterion)
accuracy = test(model, test_loader, criterion)
print('current epoch: {0}, loss: {1}, accuracy: {2}'.format(epoch, loss, accuracy))
if accuracy > best_accuracy:
best_accuracy = accuracy
# state dict of weights and masks
best_state_dict = copy.deepcopy(model.state_dict())
print('prune iteration: {0}, loss: {1}, accuracy: {2}'.format(i, loss, accuracy))
if best_accuracy > orig_accuracy:
# load weights and masks
pruner.bound_model.load_state_dict(best_state_dict)
# reset weights to original untrained model and keep masks unchanged to export winning ticket
pruner.load_model_state_dict(orig_state)
pruner.export_model('model_winning_ticket.pth', 'mask_winning_ticket.pth')
print('winning ticket has been saved: model_winning_ticket.pth, mask_winning_ticket.pth')
else:
print('winning ticket is not found in this run, you can run it again.')
examples/model_compress/pruning/mobilenetv2_end2end/Compressing MobileNetV2 with NNI Pruners.ipynb 0 → 100644
View file @ 1011377c
This source diff could not be displayed because it is too large. You can view the blob instead.
examples/model_compress/pruning/mobilenetv2_end2end/prepare_data.sh 0 → 100755
View file @ 1011377c
#!/bin/bash
# download and preprocess the Stanford Dogs dataset
mkdir -p data/stanford-dogs
# download raw data (images, annotations, and train-test split)
cd data/stanford-dogs
if [ ! -d './Images' ] ; then
if [ ! -f 'images.tar' ] ; then
wget http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar
fi
tar -xvf images.tar
fi
if [ ! -d './Annotation' ] ; then
if [ ! -f 'annotation.tar' ] ; then
wget http://vision.stanford.edu/aditya86/ImageNetDogs/annotation.tar
fi
tar -xvf annotation.tar
fi
if [ ! -f 'lists.tar' ] ; then
wget http://vision.stanford.edu/aditya86/ImageNetDogs/lists.tar
fi
tar -xvf lists.tar
cd ../..
# preprocess: train-valid-test splitting and image cropping
python preprocess.py
examples/model_compress/pruning/mobilenetv2_end2end/preprocess.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import xml.etree.ElementTree
from PIL import Image
import numpy as np
from sklearn.model_selection import train_test_split
from scipy import io
ROOT_DIR = './data/stanford-dogs/'
NUM_CATEGORIES = 120
OUT_IMAGE_SIZE = (224, 224)
RANDOM_SEED = 42 # for splitting train and validation
TRAIN_RATIO = 0.9 # train / (train + validation)
def get_bounding_box(annotation_file):
"""
Parse the annotation file and returns the bounding box information
Parameters
----------
annotation_file: path to the annotation XML file
Returns
-------
A dict containing bounding box information
"""
ret = {}
xml_root = xml.etree.ElementTree.parse(annotation_file).getroot()
bounding_box = xml_root.findall('object')[0].findall('bndbox')[0]
ret['X_min'] = int(bounding_box.findall('xmin')[0].text)
ret['X_max'] = int(bounding_box.findall('xmax')[0].text)
ret['Y_min'] = int(bounding_box.findall('ymin')[0].text)
ret['Y_max'] = int(bounding_box.findall('ymax')[0].text)
return ret
def main(root_dir):
try:
os.mkdir(root_dir + 'Processed')
os.mkdir(root_dir + 'Processed/train')
os.mkdir(root_dir + 'Processed/valid')
os.mkdir(root_dir + 'Processed/test')
except:
print('Directory already exists. Nothing done.')
exit()
# load train test splits
train_metadata = io.loadmat(root_dir + 'train_list.mat')
train_valid_file_list = [x[0][0] for x in train_metadata['file_list']]
train_valid_annotation_list = [x[0][0] for x in train_metadata['annotation_list']]
train_valid_labels = [x[0] - 1 for x in train_metadata['labels']]
train_valid_lists = [x for x in zip(train_valid_file_list, train_valid_annotation_list, train_valid_labels)]
train_lists, valid_lists = train_test_split(train_valid_lists, train_size=TRAIN_RATIO, random_state=RANDOM_SEED)
train_file_list, train_annotation_list, train_labels = zip(*train_lists)
valid_file_list, valid_annotation_list, valid_labels = zip(*valid_lists)
test_metadata = io.loadmat(root_dir + 'test_list.mat')
test_file_list = [x[0][0] for x in test_metadata['file_list']]
test_annotation_list = [x[0][0] for x in test_metadata['annotation_list']]
test_labels = [x[0] - 1 for x in test_metadata['labels']]
label2idx = {}
for split, file_list, annotation_list, labels in zip(['train', 'valid', 'test'],
[train_file_list, valid_file_list, test_file_list],
[train_annotation_list, valid_annotation_list, test_annotation_list],
[train_labels, valid_labels, test_labels]):
print('Preprocessing {} set: {} cases'.format(split, len(file_list)))
for cur_file, cur_annotation, cur_label in zip(file_list, annotation_list, labels):
label_name = cur_file.split('/')[0].split('-')[-1].lower()
if label_name not in label2idx:
label2idx[label_name] = cur_label
image = Image.open(root_dir + '/Images/' + cur_file)
# cropping and reshape
annotation_file = root_dir + '/Annotation/' + cur_annotation
bounding_box = get_bounding_box(annotation_file)
image = image.crop([bounding_box['X_min'], bounding_box['Y_min'],
bounding_box['X_max'], bounding_box['Y_max']])
image = image.convert('RGB')
image = image.resize(OUT_IMAGE_SIZE)
# Normalize and save the instance
X = np.array(image)
X = (X - np.mean(X, axis=(0, 1))) / np.std(X, axis=(0, 1)) # normalize each channel separately
# image.save(root_dir + 'Processed/' + split + '/' + image_name)
np.save(root_dir + 'Processed/' + split + '/' + cur_file.split('/')[-1].replace('.jpg', '.npy'),
{'input': X, 'label': cur_label})
# save mapping from label name to index to a dict
with open(ROOT_DIR + '/category_dict.tsv', 'w') as dict_f:
final_dict_list = sorted(list(label2idx.items()), key=(lambda x: x[-1]))
for label, index in final_dict_list:
dict_f.write('{}\t{}\n'.format(index, label))
print(final_dict_list)
if __name__ == '__main__':
main(ROOT_DIR)
examples/model_compress/pruning/mobilenetv2_end2end/pretrain.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import argparse
from time import gmtime, strftime
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
from utils import *
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def run_validation(model, valid_dataloader):
model.eval()
loss_func = nn.CrossEntropyLoss()
acc_list, loss_list = [], []
with torch.no_grad():
for i, (inputs, labels) in enumerate(tqdm(valid_dataloader)):
inputs, labels = inputs.float().to(device), labels.to(device)
preds= model(inputs)
pred_idx = preds.max(1).indices
acc = (pred_idx == labels).sum().item() / labels.size(0)
acc_list.append(acc)
loss = loss_func(preds, labels).item()
loss_list.append(loss)
valid_loss = np.array(loss_list).mean()
valid_acc = np.array(acc_list).mean()
return valid_loss, valid_acc
def run_pretrain(args):
print(args)
torch.set_num_threads(args.n_workers)
model_type = 'mobilenet_v2_torchhub'
pretrained = True # load imagenet weight
experiment_dir = 'pretrained_{}'.format(model_type) if args.experiment_dir is None else args.experiment_dir
os.mkdir(experiment_dir)
checkpoint = None
input_size = 224
n_classes = 120
log = open(experiment_dir + '/pretrain.log', 'w')
model = create_model(model_type=model_type, pretrained=pretrained, n_classes=n_classes,
input_size=input_size, checkpoint=checkpoint)
model = model.to(device)
print(model)
# count_flops(model, device=device)
train_dataset = TrainDataset('./data/stanford-dogs/Processed/train')
train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
valid_dataset = EvalDataset('./data/stanford-dogs/Processed/valid')
valid_dataloader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay)
best_valid_acc = 0.0
for epoch in range(args.n_epochs):
print('Start training epoch {}'.format(epoch))
loss_list = []
# train
model.train()
for i, (inputs, labels) in enumerate(tqdm(train_dataloader)):
optimizer.zero_grad()
inputs, labels = inputs.float().to(device), labels.to(device)
preds = model(inputs)
loss = criterion(preds, labels)
loss_list.append(loss.item())
loss.backward()
optimizer.step()
# validation
valid_loss, valid_acc = run_validation(model, valid_dataloader)
train_loss = np.array(loss_list).mean()
print('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}'.format
(epoch, train_loss, valid_loss, valid_acc))
log.write('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}\n'.format
(epoch, train_loss, valid_loss, valid_acc))
# save
if valid_acc > best_valid_acc:
best_valid_acc = valid_acc
torch.save(model.state_dict(), experiment_dir + '/checkpoint_best.pt')
log.close()
def parse_args():
parser = argparse.ArgumentParser(description='Example code for pruning MobileNetV2')
parser.add_argument('--experiment_dir', type=str, default=None,
help='directory containing the pretrained model')
parser.add_argument('--checkpoint_name', type=str, default='checkpoint_best.pt',
help='checkpoint of the pretrained model')
# finetuning parameters
parser.add_argument('--n_workers', type=int, default=16,
help='number of threads')
parser.add_argument('--n_epochs', type=int, default=180,
help='number of epochs to train the model')
parser.add_argument('--learning_rate', type=float, default=1e-4)
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument('--batch_size', type=int, default=32,
help='input batch size for training and inference')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
run_pretrain(args)
examples/model_compress/pruning/mobilenetv2_end2end/pruning_experiments.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import argparse
import copy
from time import gmtime, strftime
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
import nni
from nni.compression.pytorch import ModelSpeedup
from nni.algorithms.compression.pytorch.pruning import (
LevelPruner,
SlimPruner,
FPGMPruner,
TaylorFOWeightFilterPruner,
L1FilterPruner,
L2FilterPruner,
AGPPruner,
ActivationMeanRankFilterPruner,
ActivationAPoZRankFilterPruner
)
from utils import *
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_type = 'mobilenet_v2_torchhub'
input_size = 224
n_classes = 120
pruner_type_to_class = {'level': LevelPruner,
'l1': L1FilterPruner,
'l2': L2FilterPruner,
'slim': SlimPruner,
'fpgm': FPGMPruner,
'taylorfo': TaylorFOWeightFilterPruner,
'agp': AGPPruner,
'mean_activation': ActivationMeanRankFilterPruner,
'apoz': ActivationAPoZRankFilterPruner}
def run_eval(model, dataloader, device):
model.eval()
loss_func = nn.CrossEntropyLoss()
acc_list, loss_list = [], []
with torch.no_grad():
for i, (inputs, labels) in enumerate(tqdm(dataloader)):
inputs, labels = inputs.float().to(device), labels.to(device)
preds= model(inputs)
pred_idx = preds.max(1).indices
acc = (pred_idx == labels).sum().item() / labels.size(0)
acc_list.append(acc)
loss = loss_func(preds, labels).item()
loss_list.append(loss)
final_loss = np.array(loss_list).mean()
final_acc = np.array(acc_list).mean()
return final_loss, final_acc
def run_finetune(model, train_dataloader, valid_dataloader, device,
n_epochs=2, learning_rate=1e-4, weight_decay=0.0, log=None):
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
best_valid_acc = 0.0
best_model = None
for epoch in range(n_epochs):
print('Start finetuning epoch {}'.format(epoch))
loss_list = []
# train
model.train()
for i, (inputs, labels) in enumerate(tqdm(train_dataloader)):
optimizer.zero_grad()
inputs, labels = inputs.float().to(device), labels.to(device)
preds = model(inputs)
loss = criterion(preds, labels)
loss_list.append(loss.item())
loss.backward()
optimizer.step()
# validation
valid_loss, valid_acc = run_eval(model, valid_dataloader, device)
train_loss = np.array(loss_list).mean()
print('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}'.format
(epoch, train_loss, valid_loss, valid_acc))
if log is not None:
log.write('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}'.format
(epoch, train_loss, valid_loss, valid_acc))
if valid_acc > best_valid_acc:
best_valid_acc = valid_acc
best_model = copy.deepcopy(model).to(device)
print("Best validation accuracy: {}".format(best_valid_acc))
if log is not None:
log.write("Best validation accuracy: {}".format(best_valid_acc))
model = best_model
return model
def run_finetune_distillation(student_model, teacher_model, train_dataloader, valid_dataloader, device,
alpha, temperature,
n_epochs=2, learning_rate=1e-4, weight_decay=0.0, log=None):
optimizer = torch.optim.Adam(student_model.parameters(), lr=learning_rate, weight_decay=weight_decay)
# optimizer = torch.optim.SGD(student_model.parameters(), lr=learning_rate, momentum=0.9)
best_valid_acc = 0.0
best_model = None
for epoch in range(n_epochs):
print('Start finetuning with distillation epoch {}'.format(epoch))
loss_list = []
# train
student_model.train()
for i, (inputs, labels) in enumerate(tqdm(train_dataloader)):
optimizer.zero_grad()
inputs, labels = inputs.float().to(device), labels.to(device)
with torch.no_grad():
teacher_preds = teacher_model(inputs)
preds = student_model(inputs)
soft_loss = nn.KLDivLoss()(F.log_softmax(preds/temperature, dim=1),
F.softmax(teacher_preds/temperature, dim=1))
hard_loss = F.cross_entropy(preds, labels)
loss = soft_loss * (alpha * temperature * temperature) + hard_loss * (1. - alpha)
loss_list.append(loss.item())
loss.backward()
optimizer.step()
# validation
valid_loss, valid_acc = run_eval(student_model, valid_dataloader, device)
train_loss = np.array(loss_list).mean()
print('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}'.format
(epoch, train_loss, valid_loss, valid_acc))
if log is not None:
log.write('Epoch {}: train loss {:.4f}, valid loss {:.4f}, valid acc {:.4f}'.format
(epoch, train_loss, valid_loss, valid_acc))
if valid_acc > best_valid_acc:
best_valid_acc = valid_acc
best_model = copy.deepcopy(student_model).to(device)
print("Best validation accuracy: {}".format(best_valid_acc))
if log is not None:
log.write("Best validation accuracy: {}".format(best_valid_acc))
student_model = best_model
return student_model
def trainer_helper(model, criterion, optimizer, dataloader, device):
print("Running trainer in tuner")
for epoch in range(1):
model.train()
for i, (inputs, labels) in enumerate(tqdm(dataloader)):
optimizer.zero_grad()
inputs, labels = inputs.float().to(device), labels.to(device)
preds = model(inputs)
loss = criterion(preds, labels)
loss.backward()
optimizer.step()
def trainer_helper_with_distillation(model, teacher_model, alpha, temperature, optimizer, dataloader, device):
print("Running trainer in tuner")
for epoch in range(1):
model.train()
for i, (inputs, labels) in enumerate(tqdm(dataloader)):
optimizer.zero_grad()
inputs, labels = inputs.float().to(device), labels.to(device)
with torch.no_grad():
teacher_preds = teacher_model(inputs)
preds = model(inputs)
soft_loss = nn.KLDivLoss()(F.log_softmax(preds/temperature, dim=1),
F.softmax(teacher_preds/temperature, dim=1))
hard_loss = F.cross_entropy(preds, labels)
loss = soft_loss * (alpha * temperature * temperature) + hard_loss * (1. - alpha)
loss.backward()
optimizer.step()
def parse_args():
parser = argparse.ArgumentParser(description='Example code for pruning MobileNetV2')
parser.add_argument('--experiment_dir', type=str, required=True,
help='directory containing the pretrained model')
parser.add_argument('--checkpoint_name', type=str, default='checkpoint_best.pt',
help='checkpoint of the pretrained model')
# pruner
parser.add_argument('--pruning_mode', type=str, default='conv1andconv2',
choices=['conv0', 'conv1', 'conv2', 'conv1andconv2', 'all'])
parser.add_argument('--sparsity', type=float, default=0.5,
help='target sparsity')
parser.add_argument('--pruner_name', type=str, default='l1',
choices=['l1', 'l2', 'slim', 'agp',
'fpgm', 'mean_activation', 'apoz', 'taylorfo'],
help='pruner to use')
# for agp only
parser.add_argument('--agp_pruning_alg', default='l1',
choices=['l1', 'l2', 'slim', 'fpgm',
'mean_activation', 'apoz', 'taylorfo'],
help='pruner to use for agp')
parser.add_argument('--agp_n_iters', type=int, default=64,
help='number of iterations for agp')
parser.add_argument('--agp_n_epochs_per_iter', type=int, default=1,
help='number of epochs per iteration for agp')
# speed-up
parser.add_argument('--speed_up', action='store_true', default=False,
help='Whether to speed-up the pruned model')
# finetuning parameters
parser.add_argument('--n_workers', type=int, default=16,
help='number of threads')
parser.add_argument('--finetune_epochs', type=int, default=180,
help='number of epochs to finetune the model')
parser.add_argument('--learning_rate', type=float, default=1e-4)
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument('--batch_size', type=int, default=32,
help='input batch size for training and inference')
parser.add_argument('--kd', action='store_true', default=False,
help='Whether to use knowledge distillation')
parser.add_argument('--alpha', type=float, default=0.99,
help='Alpha for knowledge distillation loss')
parser.add_argument('--temp', type=float, default=8,
help='Temperature for knowledge distillation loss')
args = parser.parse_args()
return args
def run_pruning(args):
print(args)
torch.set_num_threads(args.n_workers)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
log = open(args.experiment_dir + '/pruning_{}_{}_sparsity{}_{}.log'.format(
args.pruner_name, args.pruning_mode, args.sparsity,
strftime("%Y%m%d%H%M", gmtime())), 'w')
train_dataset = TrainDataset('./data/stanford-dogs/Processed/train')
train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
train_dataset_for_pruner = EvalDataset('./data/stanford-dogs/Processed/train')
train_dataloader_for_pruner = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False)
valid_dataset = EvalDataset('./data/stanford-dogs/Processed/valid')
valid_dataloader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False)
test_dataset = EvalDataset('./data/stanford-dogs/Processed/test')
test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
model = create_model(model_type=model_type, pretrained=False, n_classes=n_classes,
input_size=input_size, checkpoint=args.experiment_dir + '/' + args.checkpoint_name)
model = model.to(device)
teacher_model = None
if args.kd:
teacher_model = copy.deepcopy(model)
# evaluation before pruning
# count_flops(model, log, device)
initial_loss, initial_acc = run_eval(model, test_dataloader, device)
print('Before Pruning:\nLoss: {}\nAccuracy: {}'.format(initial_loss, initial_acc))
log.write('Before Pruning:\nLoss: {}\nAccuracy: {}\n'.format(initial_loss, initial_acc))
# set up config list and pruner
config_list = []
if 'conv0' in args.pruning_mode or args.pruning_mode == 'all':
if args.pruner_name == 'slim' or (args.pruner_name == 'agp' and args.agp_pruning_alg == 'slim'):
config_list.append({
'op_names': ['features.{}.conv.0.1'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
else:
config_list.append({
'op_names': ['features.{}.conv.0.0'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
if 'conv1' in args.pruning_mode or args.pruning_mode == 'all':
if args.pruner_name == 'slim' or (args.pruner_name == 'agp' and args.agp_pruning_alg == 'slim'):
config_list.append({
'op_names': ['features.{}.conv.1.1'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
else:
config_list.append({
'op_names': ['features.{}.conv.1.0'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
if 'conv2' in args.pruning_mode or args.pruning_mode == 'all':
if args.pruner_name == 'slim' or (args.pruner_name == 'agp' and args.agp_pruning_alg == 'slim'):
config_list.append({
'op_names': ['features.{}.conv.3'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
else:
config_list.append({
'op_names': ['features.{}.conv.2'.format(x) for x in range(2, 18)],
'sparsity': args.sparsity
})
print(config_list)
kwargs = {}
if args.pruner_name in ['slim', 'taylorfo', 'mean_activation', 'apoz', 'agp']:
def trainer(model, optimizer, criterion, epoch):
if not args.kd:
return trainer_helper(model, criterion, optimizer, train_dataloader, device)
else:
return trainer_helper_with_distillation(model, teacher_model, args.alpha, args.temp, optimizer, train_dataloader, device)
kwargs = {
'trainer': trainer,
'optimizer': torch.optim.Adam(model.parameters()),
'criterion': nn.CrossEntropyLoss()
}
if args.pruner_name == 'agp':
kwargs['pruning_algorithm'] = args.agp_pruning_alg
kwargs['num_iterations'] = args.agp_n_iters
kwargs['epochs_per_iteration'] = args.agp_n_epochs_per_iter
if args.pruner_name == 'slim':
kwargs['sparsifying_training_epochs'] = 10
# pruning
pruner = pruner_type_to_class[args.pruner_name](model, config_list, **kwargs)
pruner.compress()
pruner.export_model(args.experiment_dir + '/model_temp.pth', args.experiment_dir + './mask_temp.pth')
# model speedup
pruner._unwrap_model()
if args.speed_up:
dummy_input = torch.rand(1,3,224,224).to(device)
ms = ModelSpeedup(model, dummy_input, args.experiment_dir + './mask_temp.pth')
ms.speedup_model()
print(model)
count_flops(model, log)
intermediate_loss, intermediate_acc = run_eval(model, test_dataloader, device)
print('Before Finetuning:\nLoss: {}\nAccuracy: {}'.format(intermediate_loss, intermediate_acc))
log.write('Before Finetuning:\nLoss: {}\nAccuracy: {}\n'.format(intermediate_loss, intermediate_acc))
# finetuning
if args.kd:
model = run_finetune_distillation(model, teacher_model, train_dataloader, valid_dataloader, device,
args.alpha, args.temp, n_epochs=args.finetune_epochs,
learning_rate=args.learning_rate, weight_decay=args.weight_decay)
else:
model = run_finetune(model, train_dataloader, valid_dataloader, device, n_epochs=args.finetune_epochs,
learning_rate=args.learning_rate, weight_decay=args.weight_decay)
# final evaluation
final_loss, final_acc = run_eval(model, test_dataloader, device)
print('After Pruning:\nLoss: {}\nAccuracy: {}'.format(final_loss, final_acc))
log.write('After Pruning:\nLoss: {}\nAccuracy: {}'.format(final_loss, final_acc))
# clean up
filePaths = [args.experiment_dir + '/model_tmp.pth', args.experiment_dir + '/mask_tmp.pth']
for f in filePaths:
if os.path.exists(f):
os.remove(f)
log.close()
if __name__ == '__main__':
args = parse_args()
run_pruning(args)
examples/model_compress/pruning/mobilenetv2_end2end/test.py 0 → 100644
View file @ 1011377c
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
from utils import create_model, EvalDataset, count_flops
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_type = 'mobilenet_v2_torchhub' # 'mobilenet_v1' 'mobilenet_v2' 'mobilenet_v2_torchhub'
pretrained = False # load imagenet weight (only for 'mobilenet_v2_torchhub')
checkpoint_dir = './pretrained_{}/'.format(model_type)
checkpoint = checkpoint_dir + '/checkpoint_best.pt' # model checkpoint produced by pretrain.py
input_size = 224
n_classes = 120
batch_size = 32
def run_test():
model = create_model(model_type=model_type, pretrained=pretrained, n_classes=n_classes,
input_size=input_size, checkpoint=checkpoint)
model = model.to(device)
print(model)
# count_flops(model, device=device)
test_dataset = EvalDataset('./data/stanford-dogs/Processed/test')
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
model.eval()
loss_func = nn.CrossEntropyLoss()
acc_list, loss_list = [], []
with torch.no_grad():
for i, (inputs, labels) in enumerate(tqdm(test_dataloader)):
inputs, labels = inputs.float().to(device), labels.to(device)
preds= model(inputs)
pred_idx = preds.max(1).indices
acc = (pred_idx == labels).sum().item() / labels.size(0)
acc_list.append(acc)
loss = loss_func(preds, labels).item()
loss_list.append(loss)
final_loss = np.array(loss_list).mean()
final_acc = np.array(acc_list).mean()
print('Test loss: {}\nTest accuracy: {}'.format(final_loss, final_acc))
if __name__ == '__main__':
run_test()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment