update wrap theseus rule

ppcls/arch/backbone/base/theseus_layer.py

@@ -33,16 +33,19 @@ class TheseusLayer(nn.Layer):
         return after_stop
 
     def update_res(self, return_patterns):
-        if not return_patterns:
+        if not return_patterns or isinstance(self, WrapLayer):
             return
         for layer_i in self._sub_layers:
-            if isinstance(self._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
-                self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], return_patterns)
             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].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) and not isinstance(self, WrapLayer):
-                        self._sub_layers[layer_i] = wrap_theseus(self._sub_layers[layer_i], return_patterns)
+                        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].register_forward_post_hook(
                             self._sub_layers[layer_i]._save_sub_res_hook)
                         self._sub_layers[layer_i].res_dict = self.res_dict
@@ -93,18 +96,24 @@ class WrapLayer(TheseusLayer):
     def forward(self, *inputs, **kwargs):
         self.sub_layer(*inputs, **kwargs)
 
+    def update_res(self, return_patterns):
+        if not return_patterns or not isinstance(self.sub_layer, (nn.Sequential, nn.LayerList)):
+            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].update_res(return_patterns)
                 
-def wrap_theseus(sub_layer, return_patterns):
-    if isinstance(sub_layer, (nn.Sequential, nn.LayerList)):
-        for layer_i in sub_layer._sub_layers:
-            if isinstance(sub_layer._sub_layers[layer_i], TheseusLayer):
-                continue
-            elif isinstance(sub_layer._sub_layers[layer_i], (nn.Sequential, nn.LayerList)):
-                wrap_theseus(sub_layer._sub_layers[layer_i], return_patterns)
-            elif isinstance(sub_layer._sub_layers[layer_i], nn.Layer):
-                layer_name = sub_layer._sub_layers[layer_i].full_name()
+            layer_name = self.sub_layer._sub_layers[layer_i].full_name()
             for return_pattern in return_patterns:
                 if re.match(return_pattern, layer_name):
-                        wrap_theseus(sub_layer._sub_layers[layer_i], return_patterns)
+                    self.sub_layer._sub_layers[layer_i].res_dict = self.res_dict
+
+            if isinstance(self.sub_layer._sub_layers[layer_i], TheseusLayer):
+                self.sub_layer._sub_layers[layer_i].update_res(return_patterns)
+
+
+def wrap_theseus(sub_layer):
     wrapped_layer = WrapLayer(sub_layer)
     return wrapped_layer