update wrap theseus rule
@@ -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
