From 6b21a264b23d85c7eccabc3d27d6124a679a252a Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Mon, 12 Aug 2019 00:10:57 +0300
Subject: [PATCH] environment.py: enforce processing of known networks only
---
examples/auto_compression/amc/environment.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/examples/auto_compression/amc/environment.py b/examples/auto_compression/amc/environment.py
index af6a5e7..01dfd0b 100755
--- a/examples/auto_compression/amc/environment.py
+++ b/examples/auto_compression/amc/environment.py
@@ -329,6 +329,7 @@ class DistillerWrapperEnvironment(gym.Env):
distiller.assign_layer_fq_names(model)
modules_list = [mod.distiller_name for mod in model.modules() if type(mod)==torch.nn.Conv2d]
msglogger.warning("Using the following layers: %s" % ", ".join(modules_list))
+ raise ValueError("The config file does not specify the modules to compress for %s" % app_args.arch)
self.net_wrapper = NetworkWrapper(model, app_args, services, modules_list, amc_cfg.pruning_pattern)
self.original_model_macs, self.original_model_size = self.net_wrapper.get_resources_requirements()
--
GitLab