Better logging of pruning decisions

yapt/pruning/struct_pruning.py

@@ -58,10 +58,12 @@ class StructuredPruner:
         conv_names = find_prunable_conv2d(self.syminf)
         if self.layers is not None:
             conv_names = filter_names_by_str_or_regex(conv_names, self.layers)
-        params_flops_before = self.get_params_flops(*args)
         # 1. Prune each prunable layer
+        params_flops_before = self.get_params_flops(*args)
+        nfilters_before = {}
         for name in conv_names:
             layer = rgetattr(self.module, name)
+            nfilters_before[name] = layer.weight.shape[0]
             filters = select_filters_by_norm(layer, self.get_prune_fraction(name))
             prune_conv2d_(layer, filters)  # inplace
             self.prune_decisions[name] = filters
@@ -72,11 +74,11 @@ class StructuredPruner:
         self.module(*args)
         # 4. Print some stats
         params_flops_after = self.get_params_flops(*args)
-        self.print_stats(params_flops_before, params_flops_after)
+        self.print_stats(nfilters_before, params_flops_before, params_flops_after)
 
-    def print_stats(self, before, after):
-        params_before, flops_before = before
-        params_after, flops_after = after
+    def print_stats(self, nfilters_before, params_flops_before, params_flops_after):
+        params_before, flops_before = params_flops_before
+        params_after, flops_after = params_flops_after
         log.info(
             "%.3fX reduction in parameters after pruning (%d -> %d)",
             params_before / params_after,
@@ -90,6 +92,13 @@ class StructuredPruner:
             flops_after,
         )
 
+        lines = []
+        for layer_name, filters in self.prune_decisions.items():
+            old_n = nfilters_before[layer_name]
+            lines.append(f"  {layer_name}: {len(filters)} pruned out of {old_n}")
+            lines.append(f"  filters removed: {filters}")
+        log.debug("Layer pruned: \n%s", "\n".join(lines))
+
     def get_params_flops(self, *args):
         summary = RunModelSummary(self.graphmodule).run(*args)
         return summary["params"].sum(), summary["flops"].sum()