From e298c027f3665d1a31c04af39e4b305e19377614 Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Wed, 28 Aug 2019 12:28:22 +0300
Subject: [PATCH] compress_classifier.py: added handling of
--save-untrained-model
This command command-line argument allows us to save the
randomly-initialized model before training (useful for lottery-ticket
method).
This commit was accidentally left out of the Lottery-Ticket Hypothesis
commit from Aug 26.
---
examples/classifier_compression/compress_classifier.py | 6 ++++++
1 file changed, 6 insertions(+)
diff --git a/examples/classifier_compression/compress_classifier.py b/examples/classifier_compression/compress_classifier.py
index c25c03f..500eeb5 100755
--- a/examples/classifier_compression/compress_classifier.py
+++ b/examples/classifier_compression/compress_classifier.py
@@ -161,6 +161,12 @@ class ClassifierCompressorSampleApp(classifier.ClassifierCompressor):
def __init__(self, args, script_dir):
super().__init__(args, script_dir)
early_exit_init(args)
+ # Save the randomly-initialized model before training (useful for lottery-ticket method)
+ if args.save_untrained_model:
+ ckpt_name = '_'.join((self.args.name or "", "untrained"))
+ apputils.save_checkpoint(0, self.args.arch, self.model,
+ name=ckpt_name, dir=msglogger.logdir)
+
def handle_subapps(self):
return handle_subapps(self.model, self.criterion, self.optimizer,
--
GitLab