[End-to-end demo] MobileNetV2 Compression (#4102)

examples/model_compress/pruning/mobilenetv2_end2end/prepare_data.sh

+#!/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

+# 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

+# 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

+# 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

+# 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()
+    

examples/model_compress/pruning/mobilenetv2_end2end/utils.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as transforms
+import numpy as np
+from nni.compression.pytorch.utils.counter import count_flops_params
+
+import sys
+sys.path.append('../../models')
+from mobilenet import MobileNet
+from mobilenet_v2 import MobileNetV2
+
+
+def create_model(model_type=None, n_classes=120, input_size=224, checkpoint=None, pretrained=False, width_mult=1.):
+    if model_type == 'mobilenet_v1':
+        model = MobileNet(n_class=n_classes, profile='normal')
+    elif model_type == 'mobilenet_v2':
+        model = MobileNetV2(n_class=n_classes, input_size=input_size, width_mult=width_mult)
+    elif model_type == 'mobilenet_v2_torchhub':
+        model = torch.hub.load('pytorch/vision:v0.8.1', 'mobilenet_v2', pretrained=pretrained)
+        # model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=pretrained)
+        feature_size = model.classifier[1].weight.data.size()[1]
+        replace_classifier = torch.nn.Linear(feature_size, n_classes)
+        model.classifier[1] = replace_classifier
+    elif model_type is None:
+        model = None
+    else:
+        raise RuntimeError('Unknown model_type.')
+
+    if checkpoint is not None:
+        model.load_state_dict(torch.load(checkpoint))
+
+    return model
+
+
+def get_dataloader(dataset_type, data_path, batch_size=32, shuffle=True):
+    assert dataset_type in ['train', 'eval']
+    if dataset_type == 'train':
+        ds = TrainDataset(data_path)
+    else:
+        ds = EvalDataset(data_path)
+    return DataLoader(ds, batch_size, shuffle=shuffle)
+
+
+class TrainDataset(Dataset):
+    def __init__(self, npy_dir):
+        self.root_dir = npy_dir
+        self.case_names = [self.root_dir + '/' + x for x in os.listdir(self.root_dir)]
+        
+        transform_set = [transforms.Lambda(lambda x: x),
+                         transforms.RandomRotation(30),
+                         transforms.ColorJitter(),
+                         transforms.RandomHorizontalFlip(p=1)]
+        self.transform = transforms.RandomChoice(transform_set)
+        
+    def __len__(self):
+        return len(self.case_names)
+
+    def __getitem__(self, index):
+        instance = np.load(self.case_names[index], allow_pickle=True).item()
+        x = instance['input'].transpose(2, 0, 1)     # (C, H, W)
+        x = torch.from_numpy(x).type(torch.float)    # convert to Tensor to use torchvision.transforms
+        x = self.transform(x)
+        return x, instance['label']
+
+
+class EvalDataset(Dataset):
+    def __init__(self, npy_dir):
+        self.root_dir = npy_dir
+        self.case_names = [self.root_dir + '/' + x for x in os.listdir(self.root_dir)]
+
+    def __len__(self):
+        return len(self.case_names)
+
+    def __getitem__(self, index):
+        instance = np.load(self.case_names[index], allow_pickle=True).item()
+        x = instance['input'].transpose(2, 0, 1)
+        x = torch.from_numpy(x).type(torch.float)
+        return x, instance['label']
+
+
+def count_flops(model, log=None, device=None):
+    dummy_input = torch.rand([1, 3, 256, 256])
+    if device is not None:
+        dummy_input = dummy_input.to(device)
+    flops, params, results = count_flops_params(model, dummy_input)
+    print(f"FLOPs: {flops}, params: {params}")
+    if log is not None:
+        log.write(f"FLOPs: {flops}, params: {params}\n")
+    return flops, params