fix module hook bug (#667)

* fix bugs

* modify module register

* modify var name

* modify unittest

Author: congee524

configs/recognition/tsm/tsm_r50_gpu_normalize_1x1x8_50e_kinetics400_rgb.py

@@ -6,6 +6,7 @@
 module_hooks = [
     dict(
         type='GPUNormalize',
+        hooked_module='backbone',
         hook_pos='forward_pre',
         input_format='NCHW',
         mean=[123.675, 116.28, 103.53],

mmaction/utils/module_hooks.py

@@ -7,15 +7,21 @@
 def register_module_hooks(Module, module_hooks_list):
     handles = []
     for module_hook_cfg in module_hooks_list:
+        hooked_module_name = module_hook_cfg.pop('hooked_module', 'backbone')
+        if not hasattr(Module, hooked_module_name):
+            raise ValueError(
+                f'{Module.__class__} has no {hooked_module_name}!')
+        hooked_module = getattr(Module, hooked_module_name)
         hook_pos = module_hook_cfg.pop('hook_pos', 'forward_pre')
+
         if hook_pos == 'forward_pre':
-            handle = Module.register_forward_pre_hook(
+            handle = hooked_module.register_forward_pre_hook(
                 build_from_cfg(module_hook_cfg, MODULE_HOOKS).hook_func())
         elif hook_pos == 'forward':
-            handle = Module.register_forward_hook(
+            handle = hooked_module.register_forward_hook(
                 build_from_cfg(module_hook_cfg, MODULE_HOOKS).hook_func())
         elif hook_pos == 'backward':
-            handle = Module.register_backward_hook(
+            handle = hooked_module.register_backward_hook(
                 build_from_cfg(module_hook_cfg, MODULE_HOOKS).hook_func())
         else:
             raise ValueError(

setup.cfg

@@ -19,6 +19,6 @@ line_length = 79
 multi_line_output = 0
 known_standard_library = pkg_resources,setuptools
 known_first_party = mmaction
-known_third_party = cv2,joblib,matplotlib,mmcv,numpy,pandas,pytest,scipy,seaborn,titlecase,torch,torchvision,tqdm
+known_third_party = cv2,joblib,matplotlib,mmcv,numpy,pandas,pytest,scipy,seaborn,titlecase,torch,tqdm
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY

tests/test_utils/test_module_hooks.py

@@ -1,10 +1,12 @@
 import copy
+import os.path as osp
 
+import mmcv
 import numpy as np
 import pytest
 import torch
-import torchvision.models as models
 
+from mmaction.models import build_recognizer
 from mmaction.utils import register_module_hooks
 from mmaction.utils.module_hooks import GPUNormalize
 
@@ -13,40 +15,56 @@ def test_register_module_hooks():
     _module_hooks = [
         dict(
             type='GPUNormalize',
+            hooked_module='backbone',
             hook_pos='forward_pre',
             input_format='NCHW',
             mean=[123.675, 116.28, 103.53],
             std=[58.395, 57.12, 57.375])
     ]
 
+    repo_dpath = osp.dirname(osp.dirname(osp.dirname(__file__)))
+    config_fpath = osp.join(repo_dpath, 'configs/_base_/models/tsm_r50.py')
+    config = mmcv.Config.fromfile(config_fpath)
+    config.model['backbone']['pretrained'] = None
+
     # case 1
     module_hooks = copy.deepcopy(_module_hooks)
     module_hooks[0]['hook_pos'] = 'forward_pre'
-    resnet = models.resnet50()
-    handles = register_module_hooks(resnet, module_hooks)
-    assert resnet._forward_pre_hooks[
+    recognizer = build_recognizer(config.model)
+    handles = register_module_hooks(recognizer, module_hooks)
+    assert recognizer.backbone._forward_pre_hooks[
         handles[0].id].__name__ == 'normalize_hook'
 
     # case 2
     module_hooks = copy.deepcopy(_module_hooks)
     module_hooks[0]['hook_pos'] = 'forward'
-    resnet = models.resnet50()
-    handles = register_module_hooks(resnet, module_hooks)
-    assert resnet._forward_hooks[handles[0].id].__name__ == 'normalize_hook'
+    recognizer = build_recognizer(config.model)
+    handles = register_module_hooks(recognizer, module_hooks)
+    assert recognizer.backbone._forward_hooks[
+        handles[0].id].__name__ == 'normalize_hook'
 
     # case 3
     module_hooks = copy.deepcopy(_module_hooks)
+    module_hooks[0]['hooked_module'] = 'cls_head'
     module_hooks[0]['hook_pos'] = 'backward'
-    resnet = models.resnet50()
-    handles = register_module_hooks(resnet, module_hooks)
-    assert resnet._backward_hooks[handles[0].id].__name__ == 'normalize_hook'
+    recognizer = build_recognizer(config.model)
+    handles = register_module_hooks(recognizer, module_hooks)
+    assert recognizer.cls_head._backward_hooks[
+        handles[0].id].__name__ == 'normalize_hook'
 
     # case 4
     module_hooks = copy.deepcopy(_module_hooks)
     module_hooks[0]['hook_pos'] = '_other_pos'
-    resnet = models.resnet50()
+    recognizer = build_recognizer(config.model)
+    with pytest.raises(ValueError):
+        handles = register_module_hooks(recognizer, module_hooks)
+
+    # case 5
+    module_hooks = copy.deepcopy(_module_hooks)
+    module_hooks[0]['hooked_module'] = '_other_module'
+    recognizer = build_recognizer(config.model)
     with pytest.raises(ValueError):
-        handles = register_module_hooks(resnet, module_hooks)
+        handles = register_module_hooks(recognizer, module_hooks)
 
 
 def test_gpu_normalize():
@@ -72,9 +90,8 @@ def check_normalize(origin_imgs, result_imgs, norm_cfg):
     assert gpu_normalize._mean.shape == (1, 3, 1, 1)
     imgs = np.random.randint(256, size=(2, 240, 320, 3), dtype=np.uint8)
     _input = (torch.tensor(imgs).permute(0, 3, 1, 2), )
-    resnet = models.resnet50()
     normalize_hook = gpu_normalize.hook_func()
-    _input = normalize_hook(resnet, _input)
+    _input = normalize_hook(torch.nn.Module, _input)
     result_imgs = np.array(_input[0].permute(0, 2, 3, 1))
     check_normalize(imgs, result_imgs, gpu_normalize_cfg)
 

tools/test.py

@@ -178,7 +178,8 @@ def main():
     model = build_model(
         cfg.model, train_cfg=None, test_cfg=cfg.get('test_cfg'))
 
-    register_module_hooks(model.backbone, cfg.module_hooks)
+    if len(cfg.module_hooks) > 0:
+        register_module_hooks(model, cfg.module_hooks)
 
     fp16_cfg = cfg.get('fp16', None)
     if fp16_cfg is not None:

tools/train.py

@@ -145,7 +145,8 @@ def main():
         train_cfg=cfg.get('train_cfg'),
         test_cfg=cfg.get('test_cfg'))
 
-    register_module_hooks(model.backbone, cfg.module_hooks)
+    if len(cfg.module_hooks) > 0:
+        register_module_hooks(model, cfg.module_hooks)
 
     if cfg.omnisource:
         # If omnisource flag is set, cfg.data.train should be a list