Add examples for 'op_partial_names' and use more proper words. (#4225)

nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py

@@ -136,6 +136,7 @@ class LevelPruner(BasicPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Operation types to prune.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         """
         super().__init__(model, config_list)
@@ -170,6 +171,7 @@ class NormPruner(BasicPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Conv2d and Linear are supported in NormPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         p
             The order of norm.
@@ -229,6 +231,7 @@ class L1NormPruner(NormPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Conv2d and Linear are supported in L1NormPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         mode
             'normal' or 'dependency_aware'.
@@ -260,6 +263,7 @@ class L2NormPruner(NormPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Conv2d and Linear are supported in L2NormPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         mode
             'normal' or 'dependency_aware'.
@@ -291,6 +295,7 @@ class FPGMPruner(BasicPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Conv2d and Linear are supported in FPGMPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         mode
             'normal' or 'dependency_aware'.
@@ -351,6 +356,7 @@ class SlimPruner(BasicPruner):
                 - max_sparsity_per_layer : Always used with total_sparsity. Limit the max sparsity of each layer.
                 - op_types : Only BatchNorm2d is supported in SlimPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         trainer
             A callable function used to train model or just inference. Take model, optimizer, criterion as input.
@@ -444,6 +450,7 @@ class ActivationPruner(BasicPruner):
                 - sparsity_per_layer : Equals to sparsity.
                 - op_types : Conv2d and Linear are supported in ActivationPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         trainer
             A callable function used to train model or just inference. Take model, optimizer, criterion as input.
@@ -564,6 +571,7 @@ class TaylorFOWeightPruner(BasicPruner):
                 - max_sparsity_per_layer : Always used with total_sparsity. Limit the max sparsity of each layer.
                 - op_types : Conv2d and Linear are supported in TaylorFOWeightPruner.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         trainer
             A callable function used to train model or just inference. Take model, optimizer, criterion as input.
@@ -683,6 +691,7 @@ class ADMMPruner(BasicPruner):
                 - rho : Penalty parameters in ADMM algorithm.
                 - op_types : Operation types to prune.
                 - op_names : Operation names to prune.
+                - op_partial_names: An auxiliary field collecting matched op_names in model, then this will convert to op_names.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
         trainer
             A callable function used to train model or just inference. Take model, optimizer, criterion as input.

nni/algorithms/compression/v2/pytorch/utils/pruning.py

@@ -13,6 +13,22 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]:
     '''
     Split the config by op_names if 'sparsity' or 'sparsity_per_layer' in config,
     and set the sub_config['total_sparsity'] = config['sparsity_per_layer'].
+    And every item in 'op_partial_names' will match corresponding 'op_names' in model,
+    then convert 'op_partial_names' to 'op_names' in config.
+
+    Example::
+        model = models.resnet18()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'op_partial_names': ['conv1']}]
+        pruner = L1NormPruner(model, config_list)
+        pruner.compress()
+        pruner.show_pruned_weights()
+
+    In this process, the config_list will implicitly convert to the following:
+
+    [{'op_types': ['Conv2d'], 'sparsity_per_layer': 0.8,
+        'op_names': ['conv1', 'layer1.0.conv1', 'layer1.1.conv1',
+        'layer2.0.conv1', 'layer2.1.conv1', 'layer3.0.conv1', 'layer3.1.conv1',
+        'layer4.0.conv1', 'layer4.1.conv1']}]
     '''
     for config in config_list:
         if 'sparsity' in config: