From 0040a973ebbb2218909c4a8547ee1af007014d62 Mon Sep 17 00:00:00 2001 From: weishengyu Date: Sat, 7 Aug 2021 23:00:32 +0800 Subject: [PATCH] add save res back --- ppcls/arch/backbone/base/theseus_layer.py | 32 +++++++++++------------ 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/ppcls/arch/backbone/base/theseus_layer.py b/ppcls/arch/backbone/base/theseus_layer.py index 8ef7913f..28620820 100644 --- a/ppcls/arch/backbone/base/theseus_layer.py +++ b/ppcls/arch/backbone/base/theseus_layer.py @@ -42,29 +42,29 @@ class TheseusLayer(nn.Layer): for layer_i in self._sub_layers: layer_name = self._sub_layers[layer_i].full_name() for return_pattern in return_layers: - if return_layers is not None and re.match(return_pattern, - layer_name): + if return_layers is not None and re.match(return_pattern, layer_name): self._sub_layers[layer_i].register_forward_post_hook( self._save_sub_res_hook) - def replace_sub(self, layer_name_pattern, replace_function, - recursive=True): - for k in self._sub_layers.keys(): - layer_name = self._sub_layers[k].full_name() + def _save_sub_res_hook(self, layer, input, output): + self.res_dict[layer.full_name()] = output + + def _disconnect_res_dict_hook(self, input, output): + self.res_dict = None + + def replace_sub(self, layer_name_pattern, replace_function, recursive=True): + for layer_i in self._sub_layers: + layer_name = self._sub_layers[layer_i].full_name() if re.match(layer_name_pattern, layer_name): - self._sub_layers[k] = replace_function(self._sub_layers[k]) + self._sub_layers[layer_i] = replace_function(self._sub_layers[layer_i]) if recursive: - if isinstance(self._sub_layers[k], TheseusLayer): - self._sub_layers[k].replace_sub( + if isinstance(self._sub_layers[layer_i], TheseusLayer): + self._sub_layers[layer_i].replace_sub( layer_name_pattern, replace_function, recursive) - elif isinstance(self._sub_layers[k], - nn.Sequential) or isinstance( - self._sub_layers[k], nn.LayerList): - for kk in self._sub_layers[k]._sub_layers.keys(): - self._sub_layers[k]._sub_layers[kk].replace_sub( + elif isinstance(self._sub_layers[layer_i], (nn.Sequential, nn.LayerList)): + for layer_j in self._sub_layers[layer_i]._sub_layers: + self._sub_layers[layer_i]._sub_layers[layer_j].replace_sub( layer_name_pattern, replace_function, recursive) - else: - pass ''' example of replace function: -- GitLab