From 1140298858e51aff8c02a53437e1de1d69731a60 Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Tue, 6 Nov 2018 18:36:42 +0200
Subject: [PATCH] compress_classifier.py: add an option to load a model in
serialized mode
By default, when we create a model we wrap it with DataParallel to benefit
from data-parallelism across GPUs (mainly for convolution layers).
But sometimes we don't want the sample application to do this: for
example when we receive a model that was trained serially.
This commit adds a new argument to the application to prevent
the use of DataParallel.
---
examples/classifier_compression/compress_classifier.py | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/examples/classifier_compression/compress_classifier.py b/examples/classifier_compression/compress_classifier.py
index 6c5ab59..bf22880 100755
--- a/examples/classifier_compression/compress_classifier.py
+++ b/examples/classifier_compression/compress_classifier.py
@@ -152,7 +152,9 @@ parser.add_argument('--earlyexit_thresholds', type=float, nargs='*', dest='early
help='List of EarlyExit thresholds (e.g. --earlyexit 1.2 0.9)')
parser.add_argument('--num-best-scores', dest='num_best_scores', default=1, type=int,
help='number of best scores to track and report (default: 1)')
-
+parser.add_argument('--load-serialized', dest='load_serialized', action='store_true', default=False,
+ help='Load a model without DataParallel wrapping it')
+
quant_group = parser.add_argument_group('Arguments controlling quantization at evaluation time'
'("post-training quantization)')
quant_group.add_argument('--quantize-eval', '--qe', action='store_true',
@@ -289,7 +291,8 @@ def main():
args.exiterrors = []
# Create the model
- model = create_model(args.pretrained, args.dataset, args.arch, device_ids=args.gpus)
+ model = create_model(args.pretrained, args.dataset, args.arch,
+ parallel=not args.load_serialized, device_ids=args.gpus)
compression_scheduler = None
# 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.
@@ -666,7 +669,7 @@ def earlyexit_validate_loss(output, target, criterion, args):
this_batch_size = target.size()[0]
earlyexit_validate_criterion = nn.CrossEntropyLoss(reduction='none').cuda()
for exitnum in range(args.num_exits):
- # calculate losses at each sample separately in the minibatch.
+ # calculate losses at each sample separately in the minibatch.
args.loss_exits[exitnum] = earlyexit_validate_criterion(output[exitnum], target)
# for batch_size > 1, we need to reduce this down to an average over the batch
args.losses_exits[exitnum].add(torch.mean(args.loss_exits[exitnum]))
--
GitLab