feature_vis: Implement forward_hook to other models

Author: synml

feature_visualizer.py

@@ -38,16 +38,34 @@ def hook(model, input, output):
 
     # 모델의 각 계층에 특징맵을 받아오는 hook을 등록
     feature_maps = {}
-    model.encode1.register_forward_hook(get_feature_maps(feature_maps, 'encode1'))
-    model.encode2.register_forward_hook(get_feature_maps(feature_maps, 'encode2'))
-    model.encode3.register_forward_hook(get_feature_maps(feature_maps, 'encode3'))
-    model.encode4.register_forward_hook(get_feature_maps(feature_maps, 'encode4'))
-    model.encode_end.register_forward_hook(get_feature_maps(feature_maps, 'encode_end'))
-    model.decode4.register_forward_hook(get_feature_maps(feature_maps, 'decode4'))
-    model.decode3.register_forward_hook(get_feature_maps(feature_maps, 'decode3'))
-    model.decode2.register_forward_hook(get_feature_maps(feature_maps, 'decode2'))
-    model.decode1.register_forward_hook(get_feature_maps(feature_maps, 'decode1'))
-    model.classifier.register_forward_hook(get_feature_maps(feature_maps, 'classifier'))
+    if config['model_name'] == 'UNet':
+        model.encode1.register_forward_hook(get_feature_maps(feature_maps, 'encode1'))
+        model.encode2.register_forward_hook(get_feature_maps(feature_maps, 'encode2'))
+        model.encode3.register_forward_hook(get_feature_maps(feature_maps, 'encode3'))
+        model.encode4.register_forward_hook(get_feature_maps(feature_maps, 'encode4'))
+        model.encode_end.register_forward_hook(get_feature_maps(feature_maps, 'encode_end'))
+        model.decode4.register_forward_hook(get_feature_maps(feature_maps, 'decode4'))
+        model.decode3.register_forward_hook(get_feature_maps(feature_maps, 'decode3'))
+        model.decode2.register_forward_hook(get_feature_maps(feature_maps, 'decode2'))
+        model.decode1.register_forward_hook(get_feature_maps(feature_maps, 'decode1'))
+        model.classifier.register_forward_hook(get_feature_maps(feature_maps, 'classifier'))
+    elif config['model_name'] == 'Proposed':
+        model.initial_conv.register_forward_hook(get_feature_maps(feature_maps, 'initial_conv'))
+        model.encode1.register_forward_hook(get_feature_maps(feature_maps, 'encode1'))
+        model.encode2.register_forward_hook(get_feature_maps(feature_maps, 'encode2'))
+        model.encode3.register_forward_hook(get_feature_maps(feature_maps, 'encode3'))
+        model.encode4.register_forward_hook(get_feature_maps(feature_maps, 'encode4'))
+        model.aspp.register_forward_hook(get_feature_maps(feature_maps, 'aspp'))
+        model.decode3.register_forward_hook(get_feature_maps(feature_maps, 'decode3'))
+        model.decode2.register_forward_hook(get_feature_maps(feature_maps, 'decode2'))
+        model.decode1.register_forward_hook(get_feature_maps(feature_maps, 'decode1'))
+        model.classifier.register_forward_hook(get_feature_maps(feature_maps, 'classifier'))
+    elif config['model_name'] == 'Backbone':
+        model.initial_conv.register_forward_hook(get_feature_maps(feature_maps, 'initial_conv'))
+        model.layer1.register_forward_hook(get_feature_maps(feature_maps, 'layer1'))
+        model.layer2.register_forward_hook(get_feature_maps(feature_maps, 'layer2'))
+        model.layer3.register_forward_hook(get_feature_maps(feature_maps, 'layer3'))
+        model.layer4.register_forward_hook(get_feature_maps(feature_maps, 'layer4'))
 
     # 예측
     with torch.no_grad():