diff --git a/ppcls/arch/__init__.py b/ppcls/arch/__init__.py index 0117def4cea0fabf559305dbc21fb6928d0d4230..657fa8231bcb7d9003dfc983c202ff56dd314599 100644 --- a/ppcls/arch/__init__.py +++ b/ppcls/arch/__init__.py @@ -23,6 +23,7 @@ from . import backbone, gears from .backbone import * from .gears import build_gear from .utils import * +from ppcls.arch.backbone.base.theseus_layer import TheseusLayer from ppcls.utils import logger from ppcls.utils.save_load import load_dygraph_pretrain diff --git a/ppcls/arch/backbone/base/theseus_layer.py b/ppcls/arch/backbone/base/theseus_layer.py index f328aed6f0a5109b9bb5e7e5ccc37438ff734d51..35eac5f083bae1a371119ccf35c390441f0d1f8e 100644 --- a/ppcls/arch/backbone/base/theseus_layer.py +++ b/ppcls/arch/backbone/base/theseus_layer.py @@ -38,24 +38,31 @@ 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 + def _return_dict_hook(self, layer, input, output): + res_dict = {"output": output} + for res_key in list(self.res_dict): + res_dict[res_key] = self.res_dict.pop(res_key) + return res_dict + 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() @@ -85,10 +92,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 +110,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 +125,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/hrnet.py b/ppcls/arch/backbone/legendary_models/hrnet.py index 51ad4e4f51b7104de30962d72849c9e032229e67..7c4898a1387ec5346ba06ac0995ac1ac17c9624e 100644 --- a/ppcls/arch/backbone/legendary_models/hrnet.py +++ b/ppcls/arch/backbone/legendary_models/hrnet.py @@ -367,7 +367,7 @@ class HRNet(TheseusLayer): model: nn.Layer. Specific HRNet model depends on args. """ - def __init__(self, width=18, has_se=False, class_num=1000): + def __init__(self, width=18, has_se=False, class_num=1000, return_patterns=None): super().__init__() self.width = width @@ -456,8 +456,11 @@ class HRNet(TheseusLayer): 2048, class_num, weight_attr=ParamAttr(initializer=Uniform(-stdv, stdv))) + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) - def forward(self, x, res_dict=None): + def forward(self, x): x = self.conv_layer1_1(x) x = self.conv_layer1_2(x) diff --git a/ppcls/arch/backbone/legendary_models/inception_v3.py b/ppcls/arch/backbone/legendary_models/inception_v3.py index b6403bbe6af0ffa15eab6a6548e30398a8054a2d..50fbcb4cbaef9bbc19ed65aeee181bfea6a599b5 100644 --- a/ppcls/arch/backbone/legendary_models/inception_v3.py +++ b/ppcls/arch/backbone/legendary_models/inception_v3.py @@ -454,7 +454,7 @@ class Inception_V3(TheseusLayer): model: nn.Layer. Specific Inception_V3 model depends on args. """ - def __init__(self, config, class_num=1000): + def __init__(self, config, class_num=1000, return_patterns=None): super().__init__() self.inception_a_list = config["inception_a"] @@ -496,6 +496,9 @@ class Inception_V3(TheseusLayer): class_num, weight_attr=ParamAttr(initializer=Uniform(-stdv, stdv)), bias_attr=ParamAttr()) + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) def forward(self, x): x = self.inception_stem(x) diff --git a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py index 3a14dc81d93d3d8e0a6bc3f21323bfe94f702d28..944bdb14610ee90444f02c2bfd2dd11fb7d48b48 100644 --- a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py +++ b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py @@ -102,7 +102,7 @@ class MobileNet(TheseusLayer): model: nn.Layer. Specific MobileNet model depends on args. """ - def __init__(self, scale=1.0, class_num=1000): + def __init__(self, scale=1.0, class_num=1000, return_patterns=None): super().__init__() self.scale = scale @@ -145,6 +145,9 @@ class MobileNet(TheseusLayer): int(1024 * scale), class_num, weight_attr=ParamAttr(initializer=KaimingNormal())) + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) def forward(self, x): x = self.conv(x) diff --git a/ppcls/arch/backbone/legendary_models/mobilenet_v3.py b/ppcls/arch/backbone/legendary_models/mobilenet_v3.py index 0cb550d44f976fee9aef678e4a743c41ccc79802..438e48a4f33d2cf5a945145e0cc2e69e9eb92f2f 100644 --- a/ppcls/arch/backbone/legendary_models/mobilenet_v3.py +++ b/ppcls/arch/backbone/legendary_models/mobilenet_v3.py @@ -142,7 +142,8 @@ class MobileNetV3(TheseusLayer): inplanes=STEM_CONV_NUMBER, class_squeeze=LAST_SECOND_CONV_LARGE, class_expand=LAST_CONV, - dropout_prob=0.2): + dropout_prob=0.2, + return_patterns=None): super().__init__() self.cfg = config @@ -199,6 +200,9 @@ class MobileNetV3(TheseusLayer): self.flatten = nn.Flatten(start_axis=1, stop_axis=-1) self.fc = Linear(self.class_expand, class_num) + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) def forward(self, x): x = self.conv(x) diff --git a/ppcls/arch/backbone/legendary_models/resnet.py b/ppcls/arch/backbone/legendary_models/resnet.py index 5417e2d970e74a690ada1c5a4f3d7fdedf238d11..4f79c0d75f06e6b7e188f9311902fe3397823c8f 100644 --- a/ppcls/arch/backbone/legendary_models/resnet.py +++ b/ppcls/arch/backbone/legendary_models/resnet.py @@ -269,7 +269,8 @@ class ResNet(TheseusLayer): class_num=1000, lr_mult_list=[1.0, 1.0, 1.0, 1.0, 1.0], data_format="NCHW", - input_image_channel=3): + input_image_channel=3, + return_patterns=None): super().__init__() self.cfg = config @@ -337,6 +338,9 @@ class ResNet(TheseusLayer): weight_attr=ParamAttr(initializer=Uniform(-stdv, stdv))) self.data_format = data_format + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) def forward(self, x): with paddle.static.amp.fp16_guard(): diff --git a/ppcls/arch/backbone/legendary_models/vgg.py b/ppcls/arch/backbone/legendary_models/vgg.py index a45637d2a6a89bad081ca4452497e540872dfafe..9b1750d54c2973cc08a3eb54575b53f61ff9891e 100644 --- a/ppcls/arch/backbone/legendary_models/vgg.py +++ b/ppcls/arch/backbone/legendary_models/vgg.py @@ -137,8 +137,11 @@ class VGGNet(TheseusLayer): self.fc1 = Linear(7 * 7 * 512, 4096) self.fc2 = Linear(4096, 4096) self.fc3 = Linear(4096, class_num) + if return_patterns is not None: + self.update_res(return_patterns) + self.register_forward_post_hook(self._return_dict_hook) - def forward(self, inputs, res_dict=None): + def forward(self, inputs): x = self.conv_block_1(inputs) x = self.conv_block_2(x) x = self.conv_block_3(x) @@ -152,9 +155,6 @@ class VGGNet(TheseusLayer): x = self.relu(x) x = self.drop(x) x = self.fc3(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/engine/engine.py b/ppcls/engine/engine.py index 00f5014e440824b434b3203f265d744ac887f760..886aa3796def3d6325120bc9b4a93d155c7cba02 100644 --- a/ppcls/engine/engine.py +++ b/ppcls/engine/engine.py @@ -346,6 +346,8 @@ class Engine(object): out = self.model(batch_tensor) if isinstance(out, list): out = out[0] + if isinstance(out, dict): + out = out["output"] result = self.postprocess_func(out, image_file_list) print(result) batch_data.clear()