m

Author: wanglouis49

pruning/utils.py

@@ -64,7 +64,9 @@ def prune_rate(model):
     total_nb_param = 0
     nb_zero_param = 0
 
-    for i, parameter in enumerate(model.parameters()):
+    layer_id = 0
+
+    for parameter in model.parameters():
 
         param_this_layer = 1
         for dim in parameter.data.size():
@@ -73,11 +75,13 @@ def prune_rate(model):
 
         # only pruning linear and conv layers
         if len(parameter.data.size()) != 1:
+            layer_id += 1
             zero_param_this_layer = \
                 np.count_nonzero(parameter.cpu().data.numpy()==0)
             nb_zero_param += zero_param_this_layer
 
-            print("{} layer {:.2f}% parameters pruned".format(
+            print("Layer {} | {} layer | {:.2f}% parameters pruned".format(
+                layer_id,
                 'Conv' if len(parameter.data.size()) == 4 else 'Linear',
                 100.*zero_param_this_layer/param_this_layer,
                 ))

weight_pruning.py

@@ -41,13 +41,13 @@
 if torch.cuda.is_available():
     print('CUDA ensabled.')
     net.cuda()
-print("Pretrained network loaded.")
+print("--- Pretrained network loaded ---")
 test(net, loader_test)
 
 # prune the weights
 masks = weight_prune(net, param['pruning_perc'])
 net.set_masks(masks)
-print("{}% parameters pruned.".format(param['pruning_perc']))
+print("--- {}% parameters pruned ---".format(param['pruning_perc']))
 test(net, loader_test)
 
 
@@ -60,7 +60,7 @@
 
 
 # Check accuracy and nonzeros weights in each layer
-print("After retraining...")
+print("--- After retraining ---")
 test(net, loader_test)
 prune_rate(net)