From bd62e39c34c967df9f23c5166a356f84ebbe53ab Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Tue, 20 Aug 2019 23:29:17 +0300
Subject: [PATCH] apputils/checkpoint.py - relax requirements when loading
thinned models
Allow loading a thinned model from a checkpoint file that doesn't
contain a CompressionSchedule. After thinning, we usually don't have,
or care about, a pruning schedule - so we don't need to force the
user to create one, just for saving in the checkpoint.
---
distiller/apputils/checkpoint.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/distiller/apputils/checkpoint.py b/distiller/apputils/checkpoint.py
index e1e710d..ce76482 100755
--- a/distiller/apputils/checkpoint.py
+++ b/distiller/apputils/checkpoint.py
@@ -147,7 +147,8 @@ def load_checkpoint(model, chkpt_file, optimizer=None, model_device=None, *,
if 'thinning_recipes' in checkpoint:
if 'compression_sched' not in checkpoint:
- raise KeyError("Found thinning_recipes key, but missing mandatory key compression_sched")
+ msglogger.warning("Found thinning_recipes key, but missing mandatory key compression_sched")
+ compression_scheduler = distiller.CompressionScheduler(model)
msglogger.info("Loaded a thinning recipe from the checkpoint")
# Cache the recipes in case we need them later
model.thinning_recipes = checkpoint['thinning_recipes']
--
GitLab