Refine pruning API reaction when obtain wrong pruner type. (#314)

* Refine pruning API reaction when obtain wrong pruner type.

Signed-off-by: YIYANGCAI <[email protected]>

* Add pattern_lock valid type.

Signed-off-by: YIYANGCAI <[email protected]>

---------

Signed-off-by: YIYANGCAI <[email protected]>

Author: YIYANGCAI

neural_compressor/pruner/criteria.py

@@ -19,14 +19,14 @@
 torch = LazyImport('torch')
 
 
-CRITERIAS = {}
+CRITERIA = {}
 
 
 def register_criterion(name):
     """Register a criterion to the registry."""
 
     def register(criterion):
-        CRITERIAS[name] = criterion
+        CRITERIA[name] = criterion
         return criterion
 
     return register
@@ -35,9 +35,9 @@ def register(criterion):
 def get_criterion(config, modules):
     """Get registered criterion class."""
     name = config["criterion_type"]
-    if name not in CRITERIAS.keys():
-        assert False, f"criteria does not support {name}, currently only support {CRITERIAS.keys()}"
-    return CRITERIAS[name](modules, config)
+    if name not in CRITERIA.keys():
+        assert False, f"criteria does not support {name}, currently only support {CRITERIA.keys()}"
+    return CRITERIA[name](modules, config)
 
 
 class PruningCriterion:

neural_compressor/pruner/pruners.py

@@ -21,13 +21,12 @@
 torch = LazyImport('torch')
 from .patterns import get_pattern
 from .schedulers import get_scheduler
-from .criteria import get_criterion, CRITERIAS
+from .criteria import get_criterion, CRITERIA
 from .regs import get_reg
 from .logger import logger
 
 PRUNERS = {}
 
-
 def register_pruner(name):
     """Class decorator to register a Pruner subclass to the registry.
 
@@ -48,6 +47,14 @@ def register(pruner):
 
     return register
 
+def parse_valid_pruner_types():
+    """Get all valid pruner names"""
+    valid_pruner_types = []
+    for x in CRITERIA.keys():
+        for p in ["", "_progressive"]:
+            valid_pruner_types.append(x + p)
+    valid_pruner_types.append("pattern_lock")
+    return valid_pruner_types
 
 def get_pruner(config, modules):
     """Get registered pruner class.
@@ -71,7 +78,7 @@ def get_pruner(config, modules):
         # if progressive, delete "progressive" words and reset config["progressive"]
         name = config["pruning_type"][0:-12]
         config["progressive"] = True
-    if name in CRITERIAS:
+    if name in CRITERIA:
         if config["progressive"] == False:
             config['criterion_type'] = name
             name = "basic"  ##return the basic pruner
@@ -80,7 +87,7 @@ def get_pruner(config, modules):
             name = "progressive"  ## return the progressive pruner
 
     if name not in PRUNERS.keys():
-        assert False, f"does not support {name}, currently only support {PRUNERS.keys()}"
+        assert False, f"does not support {name}, currently only support {parse_valid_pruner_types()}"
     return PRUNERS[name](config, modules)