mirrors
/
PaddleClas
mirror of https://github.com/PaddlePaddle/PaddleClas.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

update wrap theseus rule

pull/1116/head
weishengyu 2021-08-08 15:55:34 +08:00
parent 1985075550
commit 2cb3bf66dd
1 changed files with 32 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
ppcls/arch/backbone/base/theseus_layer.py
View File

@ -33,21 +33,24 @@ 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()
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)
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):
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):
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
if isinstance(self._sub_layers[layer_i], TheseusLayer):
self._sub_layers[layer_i].update_res(return_patterns)
def _save_sub_res_hook(self, layer, input, output):
if self.res_dict is not None:
@ -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)
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
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()
for return_pattern in return_patterns:
if re.match(return_pattern, layer_name):
wrap_theseus(sub_layer._sub_layers[layer_i], return_patterns)
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