Supported add hook for root module in intermediate layers KD (#1029)

Signed-off-by: Ye, Xinyu <[email protected]>
intel · Jun 27, 2023 · 13aec25 · 13aec25
1 parent da99d28
commit 13aec25
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/neural_compressor/compression/distillation/criterions.py b/neural_compressor/compression/distillation/criterions.py
@@ -881,13 +881,14 @@ def register_hooks_for_models(self):
         """
         from neural_compressor.experimental.common import torch_utils
         def register_model_forward_hook(model, path, output_process='', student=False):
-            nodes = path.split('.')
             module = model
-            for node in nodes:
-                try:
-                    module = module.__getattr__(node)
-                except:
-                    raise AttributeError('There is no path {} in the model.'.format(path))
+            if path != '':
+                nodes = path.split('.')
+                for node in nodes:
+                    try:
+                        module = module.__getattr__(node)
+                    except:
+                        raise AttributeError('There is no path {} in the model.'.format(path))
             return module.register_forward_hook(
                 torch_utils.get_activation(path, output_process, student)
             )

diff --git a/test/distillation/test_distillation_2.x.py b/test/distillation/test_distillation_2.x.py
@@ -72,12 +72,15 @@ def test_distillation(self):
     def test_distillation_intermediate_layers(self):
         criterion = nn.CrossEntropyLoss()
         distillation_criterion_conf = IntermediateLayersKnowledgeDistillationLossConfig(
-            layer_mappings=[['layer1.0', ],
-                            [['layer1.1.conv1', ''], ['layer1.1.conv1', '0']],],
-            loss_types=['KL', 'MSE'],
-            loss_weights=[0.5, 0.5])
-
-        distillation_criterion_conf.config.IntermediateLayersKnowledgeDistillationLoss.layer_mappings[1][1][-1] = \
+            layer_mappings=[
+                ['', ],
+                ['layer1.0', ],
+                [['layer1.1.conv1', ''], ['layer1.1.conv1', '0']],
+            ],
+            loss_types=['L1', 'KL', 'MSE'],
+            loss_weights=[0.5, 0.2, 0.3])
+
+        distillation_criterion_conf.config.IntermediateLayersKnowledgeDistillationLoss.layer_mappings[2][1][-1] = \
                 lambda x: x[:, :2,...]
         optimizer = torch.optim.SGD(self.student_model.parameters(), lr=0.0001)
         conf = DistillationConfig(self.teacher_model, distillation_criterion_conf)