Carica un nuovo file

Reference Software/UFV1.0-Pruning/src/pruners/__init__.py

+import torch.nn as nn
+import torch_pruning as tp
+from functools import partial
+
+# Special Pruners
+from pruners.UpsamplePruner import UpsamplePruner
+
+def pruner_constructor(args, model, train_data, device, pruning_rateo = None):
+    sparsity_learning = False
+
+    # Random 
+    if args.pruning_method == "random":
+        imp = tp.importance.RandomImportance()
+        pruner_entry = partial(tp.pruner.MagnitudePruner)
+    # Greg
+    elif args.pruning_method == "growing_reg":
+         sparsity_learning = True
+        imp = tp.importance.GroupNormImportance(p=2, group_reduction=False)
+        pruner_entry = partial(tp.pruner.GrowingRegPruner, reg=args.reg, delta_reg=args.delta_reg)
+
+    else:
+        raise Exception("Invalid Pruning Method:", args.pruning_method)
+    
+
+    num_heads = {}
+    ignored_layers = []
+    # Cycle to extract key modules from the model
+   
+    # TODO: Add layer to be ignored for pruning (usually first and last layers)
+    # for m in model.modules():
+       
+        #ignored_layers.append(m)
+    
+        # Ignore-By-Parameters
+        # Add Upsample blocks to ignore (we eill add later them to a separate pruning function)
+        # ignored_layers.append(m)
+            
+    
+    # Special Custom Pruners maps
+    # TODO: Upsample and UpsampleOneStep are not imported
+    custom_pruners = {
+        #E.G., UpsapleBlockTypeOfYourModel: UpsamplePruner(), 
+    }
+           
+
+    pruning_ratio_dict = {}
+    pruner = pruner_entry(
+        model,
+        importance=imp,
+        example_inputs=train_data.to(device),
+        iterative_steps=args.pruning_steps,
+        pruning_ratio=pruning_rateo if pruning_rateo else args.pruning_target_ratio,
+        pruning_ratio_dict=pruning_ratio_dict,
+        max_pruning_ratio=args.max_pruning_ratio,
+        global_pruning=False,
+        num_heads=num_heads,
+        ignored_layers=ignored_layers,
+        customized_pruners=custom_pruners,
+        root_module_types=[ nn.modules.conv._ConvNd, nn.Linear ]
+    )
+    return pruner
\ No newline at end of file