From 0036011d2141c665ae5152283dcd24ae9a08fa34 Mon Sep 17 00:00:00 2001
From: Neta Zmora <31280975+nzmora@users.noreply.github.com>
Date: Mon, 23 Sep 2019 14:56:59 +0300
Subject: [PATCH] User-registered model (#391)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Add a jupyter notebook showing how to register a user's (external) image-classification model.
Contains fixes to the previous models extension mechanism, and relaxation of the `args' requirements in apputils/image_classifier.py.
apputils/image_classifier.py –
*when self.logdir is None:
-use NullLogger
-skip save_checkpoint
*return training log from run_training_loop()
*don’t log if script_dir or output_dir are not set.
*Fix params_nnz_cnt in update_training_scores_history()
data_loggers/logger.py – add NullLogger which does not log
---
distiller/apputils/image_classifier.py | 45 ++++---
distiller/data_loggers/__init__.py | 2 +-
distiller/data_loggers/logger.py | 5 +-
distiller/models/__init__.py | 10 +-
jupyter/tutorial__adding_new_model.ipynb | 158 +++++++++++++++++++++++
5 files changed, 196 insertions(+), 24 deletions(-)
create mode 100644 jupyter/tutorial__adding_new_model.ipynb
diff --git a/distiller/apputils/image_classifier.py b/distiller/apputils/image_classifier.py
index b008746..23c07a4 100755
--- a/distiller/apputils/image_classifier.py
+++ b/distiller/apputils/image_classifier.py
@@ -62,8 +62,11 @@ class ClassifierCompressor(object):
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
- self.tflogger = TensorBoardLogger(msglogger.logdir)
- self.pylogger = PythonLogger(msglogger)
+ if not self.logdir:
+ self.pylogger = self.tflogger = NullLogger()
+ else:
+ self.tflogger = TensorBoardLogger(msglogger.logdir)
+ self.pylogger = PythonLogger(msglogger)
(self.model, self.compression_scheduler, self.optimizer,
self.start_epoch, self.ending_epoch) = _init_learner(args)
@@ -133,15 +136,16 @@ class ClassifierCompressor(object):
def _finalize_epoch(self, epoch, perf_scores_history, top1, top5):
# Update the list of top scores achieved so far, and save the checkpoint
- update_training_scores_history(perf_scores_history, self.model,
+ update_training_scores_history(perf_scores_history, self.model,
top1, top5, epoch, self.args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {'current_top1': top1,
'best_top1': perf_scores_history[0].top1,
'best_epoch': perf_scores_history[0].epoch}
- apputils.save_checkpoint(epoch, self.args.arch, self.model, optimizer=self.optimizer,
- scheduler=self.compression_scheduler, extras=checkpoint_extras,
- is_best=is_best, name=self.args.name, dir=msglogger.logdir)
+ if msglogger.logdir:
+ apputils.save_checkpoint(epoch, self.args.arch, self.model, optimizer=self.optimizer,
+ scheduler=self.compression_scheduler, extras=checkpoint_extras,
+ is_best=is_best, name=self.args.name, dir=msglogger.logdir)
def run_training_loop(self):
@@ -152,12 +156,6 @@ class ClassifierCompressor(object):
validate_one_epoch
finalize_epoch
"""
- if self.start_epoch >= self.ending_epoch:
- msglogger.error(
- 'epoch count is too low, starting epoch is {} but total epochs set to {}'.format(
- self.start_epoch, self.ending_epoch))
- raise ValueError('Epochs parameter is too low. Nothing to do.')
-
# Load the datasets lazily
self.load_datasets()
@@ -166,6 +164,7 @@ class ClassifierCompressor(object):
msglogger.info('\n')
top1, top5, loss = self.train_validate_with_scheduling(epoch)
self._finalize_epoch(epoch, perf_scores_history, top1, top5)
+ return perf_scores_history
def validate(self, epoch=-1):
self.load_datasets()
@@ -231,7 +230,7 @@ def init_classifier_compression_arg_parser():
help='collect activation statistics on phases: train, valid, and/or test'
' (WARNING: this slows down training)')
parser.add_argument('--activation-histograms', '--act-hist',
- type=float_range(exc_min=True),
+ type=distiller.utils.float_range_argparse_checker(exc_min=True),
metavar='PORTION_OF_TEST_SET',
help='Run the model in evaluation mode on the specified portion of the test dataset and '
'generate activation histograms. NOTE: This slows down evaluation significantly')
@@ -252,6 +251,8 @@ def init_classifier_compression_arg_parser():
help='an optional parameter for sensitivity testing '
'providing the range of sparsities to test.\n'
'This is equivalent to creating sensitivities = np.arange(start, stop, step)')
+ parser.add_argument('--extras', default=None, type=str,
+ help='file with extra configuration information')
parser.add_argument('--deterministic', '--det', action='store_true',
help='Ensure deterministic execution for re-producible results.')
parser.add_argument('--seed', type=int, default=None,
@@ -291,9 +292,11 @@ def init_classifier_compression_arg_parser():
def _init_logger(args, script_dir):
- module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
-
+ if not script_dir or not hasattr(args, "output_dir") or not args.output_dir:
+ msglogger.logdir = None
+ return None
+ module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'),
@@ -307,6 +310,7 @@ def _init_logger(args, script_dir):
msglogger.debug("Distiller: %s", distiller.__version__)
return msglogger.logdir
+
def _config_determinism(args):
if args.evaluate:
args.deterministic = True
@@ -328,6 +332,7 @@ def _config_determinism(args):
cudnn.benchmark = True
msglogger.info("Random seed: %d", args.seed)
+
def _config_compute_device(args):
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
@@ -400,7 +405,13 @@ def _init_learner(args):
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
- return model, compression_scheduler, optimizer, start_epoch, args.epochs
+ ending_epoch = args.epochs
+ if start_epoch >= ending_epoch:
+ msglogger.error(
+ 'epoch count is too low, starting epoch is {} but total epochs set to {}'.format(
+ start_epoch, ending_epoch))
+ raise ValueError('Epochs parameter is too low. Nothing to do.')
+ return model, compression_scheduler, optimizer, start_epoch, ending_epoch
def create_activation_stats_collectors(model, *phases):
@@ -725,7 +736,7 @@ def update_training_scores_history(perf_scores_history, model, top1, top5, epoch
perf_scores_history.sort(key=operator.attrgetter('params_nnz_cnt', 'top1', 'top5', 'epoch'), reverse=True)
for score in perf_scores_history[:num_best_scores]:
msglogger.info('==> Best [Top1: %.3f Top5: %.3f Sparsity:%.2f Params: %d on epoch: %d]',
- score.top1, score.top5, score.sparsity, -score.params_nnz_cnt, score.epoch)
+ score.top1, score.top5, score.sparsity, score.params_nnz_cnt, score.epoch)
def earlyexit_loss(output, target, criterion, args):
diff --git a/distiller/data_loggers/__init__.py b/distiller/data_loggers/__init__.py
index a3351f0..f40d718 100755
--- a/distiller/data_loggers/__init__.py
+++ b/distiller/data_loggers/__init__.py
@@ -15,7 +15,7 @@
#
from .collector import *
-from .logger import PythonLogger, TensorBoardLogger, CsvLogger
+from .logger import PythonLogger, TensorBoardLogger, CsvLogger, NullLogger
del logger
del collector
diff --git a/distiller/data_loggers/logger.py b/distiller/data_loggers/logger.py
index bc99a0d..6a4206d 100755
--- a/distiller/data_loggers/logger.py
+++ b/distiller/data_loggers/logger.py
@@ -36,7 +36,7 @@ from contextlib import ExitStack
import os
#msglogger = logging.getLogger()
-__all__ = ['PythonLogger', 'TensorBoardLogger', 'CsvLogger']
+__all__ = ['PythonLogger', 'TensorBoardLogger', 'CsvLogger', 'NullLogger']
class DataLogger(object):
@@ -64,6 +64,9 @@ class DataLogger(object):
def log_model_buffers(self, model, buffer_names, tag_prefix, epoch, completed, total, freq):
pass
+# Log to null-space
+NullLogger = DataLogger
+
class PythonLogger(DataLogger):
def __init__(self, logger):
diff --git a/distiller/models/__init__.py b/distiller/models/__init__.py
index 83e6539..95edfcb 100755
--- a/distiller/models/__init__.py
+++ b/distiller/models/__init__.py
@@ -103,9 +103,9 @@ def create_model(pretrained, dataset, arch, parallel=True, device_ids=None):
model = _create_mnist_model(arch, pretrained)
except ValueError:
if _is_registered_extension(arch, dataset, pretrained):
- model = _get_extension_model(arch, dataset)
+ model = _create_extension_model(arch, dataset)
else:
- raise ValueError('Could not recognize dataset {} and model {} pair'.format(dataset, arch))
+ raise ValueError('Could not recognize dataset {} and arch {} pair'.format(dataset, arch))
msglogger.info("=> created a %s%s model with the %s dataset" % ('pretrained ' if pretrained else '',
arch, dataset))
@@ -200,10 +200,10 @@ def register_user_model(arch, dataset, model):
def _is_registered_extension(arch, dataset, pretrained):
try:
- return _model_extensions[(arch, dataset)]
+ return _model_extensions[(arch, dataset)] is not None
except KeyError:
return None
-def _get_extension_model(arch, dataset):
- return _model_extensions[(arch, dataset)]
\ No newline at end of file
+def _create_extension_model(arch, dataset):
+ return _model_extensions[(arch, dataset)]()
\ No newline at end of file
diff --git a/jupyter/tutorial__adding_new_model.ipynb b/jupyter/tutorial__adding_new_model.ipynb
new file mode 100644
index 0000000..a1a467a
--- /dev/null
+++ b/jupyter/tutorial__adding_new_model.ipynb
@@ -0,0 +1,158 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Tutorial: Adding a new image-classification model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "import distiller\n",
+ "import torch.nn as nn\n",
+ "from distiller.models import register_user_model\n",
+ "import distiller.apputils.image_classifier as classifier"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "class MyModel(nn.Module): \n",
+ " def __init__(self):\n",
+ " super().__init__()\n",
+ " self.conv1 = nn.Conv2d(1, 20, 5, 1)\n",
+ " self.relu1 = nn.ReLU(inplace=False)\n",
+ " self.pool1 = nn.MaxPool2d(2, 2)\n",
+ " self.conv2 = nn.Conv2d(20, 50, 5, 1)\n",
+ " self.relu2 = nn.ReLU(inplace=False)\n",
+ " self.pool2 = nn.MaxPool2d(2, 2)\n",
+ " self.avgpool = nn.AvgPool2d(4, stride=1)\n",
+ " self.fc = nn.Linear(50, 10)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " x = self.pool1(self.relu1(self.conv1(x)))\n",
+ " x = self.pool2(self.relu2(self.conv2(x)))\n",
+ " x = self.avgpool(x)\n",
+ " x = x.view(x.size(0), -1)\n",
+ " x = self.fc(x)\n",
+ " return x\n",
+ "\n",
+ "def my_model():\n",
+ " return MyModel()\n",
+ " \n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[{'params_nnz_cnt': -26000.0, 'sparsity': 0.0, 'top1': 92.80000000000001, 'top5': 99.63333333333333, 'epoch': 0}]\n"
+ ]
+ }
+ ],
+ "source": [
+ "distiller.models.register_user_model(arch=\"MyModel\", dataset=\"mnist\", model=my_model)\n",
+ "model = distiller.models.create_model(pretrained=True, dataset=\"mnist\", arch=\"MyModel\")\n",
+ "assert model is not None\n",
+ "\n",
+ "\n",
+ "def init_jupyter_default_args(args):\n",
+ " args.output_dir = None\n",
+ " args.evaluate = False\n",
+ " args.seed = None\n",
+ " args.deterministic = False\n",
+ " args.cpu = True\n",
+ " args.gpus = None\n",
+ " args.load_serialized = False\n",
+ " args.deprecated_resume = None\n",
+ " args.resumed_checkpoint_path = None\n",
+ " args.load_model_path = None\n",
+ " args.reset_optimizer = False\n",
+ " args.lr = args.momentum = args.weight_decay = 0.\n",
+ " args.compress = None\n",
+ " args.epochs = 0\n",
+ " args.activation_stats = list()\n",
+ " args.batch_size = 1\n",
+ " args.workers = 1\n",
+ " args.validation_split = 0.1\n",
+ " args.effective_train_size = args.effective_valid_size = args.effective_test_size = 1.\n",
+ " args.log_params_histograms = False\n",
+ " args.print_freq = 1\n",
+ " args.masks_sparsity = False\n",
+ " args.display_confusion = False\n",
+ " args.num_best_scores = 1\n",
+ " args.name = \"\"\n",
+ "\n",
+ "\n",
+ "def config_learner_args(args, arch, dataset, dataset_path, pretrained, sgd_args, batch, epochs):\n",
+ " args.arch = \"MyModel\"\n",
+ " args.dataset = \"mnist\"\n",
+ " args.data = \"/datasets/mnist/\"\n",
+ " args.pretrained = False\n",
+ " args.lr = sgd_args[0]\n",
+ " args.momentum = sgd_args[1]\n",
+ " args.weight_decay = sgd_args[2]\n",
+ " args.batch_size = 256\n",
+ " args.epochs = epochs\n",
+ "\n",
+ "args = classifier.init_classifier_compression_arg_parser()\n",
+ "init_jupyter_default_args(args)\n",
+ "config_learner_args(args, \"MyModel\", \"mnist\", \"/datasets/mnist/\", False, (0.1, 0.9, 1e-4) , 256, 1)\n",
+ "app = classifier.ClassifierCompressor(args, script_dir=os.path.dirname(\".\"))\n",
+ "\n",
+ "# Run the training loop\n",
+ "perf_scores_history = app.run_training_loop()\n",
+ "print(perf_scores_history)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.6.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
--
GitLab