[End-to-end demo] MobileNetV2 Compression

examples/model_compress/pruning/mobilenetv2_end2end/prepare_data.sh

@@ -0,0 +1,32 @@
+#!/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,0 +1,103 @@
+# 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,0 +1,123 @@
+# 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)