From 5131956d15d98fa3964e8b41efdf485cefd0bc88 Mon Sep 17 00:00:00 2001
From: weishengyu <weishengyu@baidu.com>
Date: Mon, 6 Sep 2021 16:18:57 +0800
Subject: [PATCH] dbg theseus

---
 ppcls/arch/backbone/base/theseus_layer.py     | 29 ++++++++++---------
 .../backbone/legendary_models/mobilenet_v1.py |  5 +++-
 ppcls/arch/backbone/legendary_models/vgg.py   |  2 +-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/ppcls/arch/backbone/base/theseus_layer.py b/ppcls/arch/backbone/base/theseus_layer.py
index f328aed6..29697e65 100644
--- a/ppcls/arch/backbone/base/theseus_layer.py
+++ b/ppcls/arch/backbone/base/theseus_layer.py
@@ -38,20 +38,21 @@ class TheseusLayer(nn.Layer):
         for layer_i in self._sub_layers:
             layer_name = self._sub_layers[layer_i].full_name()
             if isinstance(self._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
-                self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i])
-                self._sub_layers[layer_i].res_dict = self.res_dict
+                self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], self.res_dict)
                 self._sub_layers[layer_i].update_res(return_patterns)
             else:
                 for return_pattern in return_patterns:
                     if re.match(return_pattern, layer_name):
                         if not isinstance(self._sub_layers[layer_i], TheseusLayer):
-                            self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i])
+                            self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], self.res_dict)
+                        else:
+                            self._sub_layers[layer_i].res_dict = self.res_dict
+
                         self._sub_layers[layer_i].register_forward_post_hook(
                             self._sub_layers[layer_i]._save_sub_res_hook)
-                        self._sub_layers[layer_i].res_dict = self.res_dict
-                if isinstance(self._sub_layers[layer_i], TheseusLayer):
-                    self._sub_layers[layer_i].res_dict = self.res_dict
-                    self._sub_layers[layer_i].update_res(return_patterns)
+            if isinstance(self._sub_layers[layer_i], TheseusLayer):
+                self._sub_layers[layer_i].res_dict = self.res_dict
+                self._sub_layers[layer_i].update_res(return_patterns)
 
     def _save_sub_res_hook(self, layer, input, output):
         self.res_dict[layer.full_name()] = output
@@ -85,10 +86,12 @@ class TheseusLayer(nn.Layer):
 
 
 class WrapLayer(TheseusLayer):
-    def __init__(self, sub_layer):
+    def __init__(self, sub_layer, res_dict=None):
         super(WrapLayer, self).__init__()
         self.sub_layer = sub_layer
         self.name = sub_layer.full_name()
+        if res_dict is not None:
+            self.res_dict = res_dict
 
     def full_name(self):
         return self.name
@@ -101,14 +104,14 @@ class WrapLayer(TheseusLayer):
             return
         for layer_i in self.sub_layer._sub_layers:
             if isinstance(self.sub_layer._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
-                self.sub_layer._sub_layers[layer_i] = wrap_theseus(self.sub_layer._sub_layers[layer_i])
-                self.sub_layer._sub_layers[layer_i].res_dict = self.res_dict
+                self.sub_layer._sub_layers[layer_i] = wrap_theseus(self.sub_layer._sub_layers[layer_i], self.res_dict)
                 self.sub_layer._sub_layers[layer_i].update_res(return_patterns)
+            elif isinstance(self.sub_layer._sub_layers[layer_i], TheseusLayer):
+                self.sub_layer._sub_layers[layer_i].res_dict = self.res_dict
 
             layer_name = self.sub_layer._sub_layers[layer_i].full_name()
             for return_pattern in return_patterns:
                 if re.match(return_pattern, layer_name):
-                    self.sub_layer._sub_layers[layer_i].res_dict = self.res_dict
                     self.sub_layer._sub_layers[layer_i].register_forward_post_hook(
                         self._sub_layers[layer_i]._save_sub_res_hook)
 
@@ -116,6 +119,6 @@ class WrapLayer(TheseusLayer):
                 self.sub_layer._sub_layers[layer_i].update_res(return_patterns)
 
 
-def wrap_theseus(sub_layer):
-    wrapped_layer = WrapLayer(sub_layer)
+def wrap_theseus(sub_layer, res_dict=None):
+    wrapped_layer = WrapLayer(sub_layer, res_dict)
     return wrapped_layer
diff --git a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
index 3a14dc81..bacac5b5 100644
--- a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
+++ b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
@@ -146,12 +146,15 @@ class MobileNet(TheseusLayer):
             class_num,
             weight_attr=ParamAttr(initializer=KaimingNormal()))
 
-    def forward(self, x):
+    def forward(self, x, res_dict=None):
         x = self.conv(x)
         x = self.blocks(x)
         x = self.avg_pool(x)
         x = self.flatten(x)
         x = self.fc(x)
+        if self.res_dict and res_dict is not None:
+            for res_key in list(self.res_dict):
+                res_dict[res_key] = self.res_dict.pop(res_key)
         return x
 
 
diff --git a/ppcls/arch/backbone/legendary_models/vgg.py b/ppcls/arch/backbone/legendary_models/vgg.py
index a45637d2..a8980929 100644
--- a/ppcls/arch/backbone/legendary_models/vgg.py
+++ b/ppcls/arch/backbone/legendary_models/vgg.py
@@ -111,7 +111,7 @@ class VGGNet(TheseusLayer):
         model: nn.Layer. Specific VGG model depends on args.
     """
 
-    def __init__(self, config, stop_grad_layers=0, class_num=1000, return_patterns=None):
+    def __init__(self, config, stop_grad_layers=0, class_num=1000):
         super().__init__()
 
         self.stop_grad_layers = stop_grad_layers
-- 
GitLab