diff --git a/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py b/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py index 750b488ff941e8b2e2624a7e704345f9ca690920..4fad976d0a27d08ce28141870b5a536143dbab13 100644 --- a/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py +++ b/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py @@ -24,7 +24,6 @@ from paddle.nn.initializer import KaimingNormal from ppcls.arch.backbone.base.theseus_layer import TheseusLayer from ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url -#TODO(gaotingquan): upload pretrained to bos MODEL_URLS = { "PPLCNetV2_base": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNetV2_base_pretrained.pdparams", @@ -252,14 +251,16 @@ class RepDepthwiseSeparable(TheseusLayer): class PPLCNetV2(TheseusLayer): def __init__(self, - scale=1.0, - depths=[2, 2, 6, 2], + scale, + depths, class_num=1000, - dropout_prob=0.2, - expand=1280): + dropout_prob=0, + use_last_conv=True, + class_expand=1280): super().__init__() self.scale = scale - self.expand = expand + self.use_last_conv = use_last_conv + self.class_expand = class_expand self.stem = nn.Sequential(* [ ConvBNLayer( @@ -273,77 +274,33 @@ class PPLCNetV2(TheseusLayer): dw_size=3) ]) - # stage1 - in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut = NET_CONFIG[ - "stage1"] - self.stage1 = nn.Sequential(* [ - RepDepthwiseSeparable( - in_channels=make_divisible((in_channels if i == 0 else - in_channels * 2) * scale), - out_channels=make_divisible(in_channels * 2 * scale), - stride=2 if i == 0 else 1, - dw_size=kernel_size, - split_pw=split_pw, - use_rep=use_rep, - use_se=use_se, - use_shortcut=use_shortcut, ) for i in range(depths[0]) - ]) - - # stage2 - in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut = NET_CONFIG[ - "stage2"] - self.stage2 = nn.Sequential(* [ - RepDepthwiseSeparable( - in_channels=make_divisible((in_channels if i == 0 else - in_channels * 2) * scale), - out_channels=make_divisible(in_channels * 2 * scale), - stride=2 if i == 0 else 1, - dw_size=kernel_size, - split_pw=split_pw, - use_rep=use_rep, - use_se=use_se, - use_shortcut=use_shortcut, ) for i in range(depths[1]) - ]) - - # stage3 - in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut = NET_CONFIG[ - "stage3"] - self.stage3 = nn.Sequential(* [ - RepDepthwiseSeparable( - in_channels=make_divisible((in_channels if i == 0 else - in_channels * 2) * scale), - out_channels=make_divisible(in_channels * 2 * scale), - stride=2 if i == 0 else 1, - dw_size=kernel_size, - split_pw=split_pw, - use_rep=use_rep, - use_se=use_se, - use_shortcut=use_shortcut, ) for i in range(depths[2]) - ]) - - # stage4 - in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut = NET_CONFIG[ - "stage4"] - self.stage4 = nn.Sequential(* [ - RepDepthwiseSeparable( - in_channels=make_divisible((in_channels if i == 0 else - in_channels * 2) * scale), - out_channels=make_divisible(in_channels * 2 * scale), - stride=2 if i == 0 else 1, - dw_size=kernel_size, - split_pw=split_pw, - use_rep=use_rep, - use_se=use_se, - use_shortcut=use_shortcut, ) for i in range(depths[3]) - ]) + # stages + self.stages = nn.LayerList() + for depth_idx, k in enumerate(NET_CONFIG): + in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut = NET_CONFIG[ + k] + self.stages.append( + nn.Sequential(* [ + RepDepthwiseSeparable( + in_channels=make_divisible((in_channels if i == 0 else + in_channels * 2) * scale), + out_channels=make_divisible(in_channels * 2 * scale), + stride=2 if i == 0 else 1, + dw_size=kernel_size, + split_pw=split_pw, + use_rep=use_rep, + use_se=use_se, + use_shortcut=use_shortcut) + for i in range(depths[depth_idx]) + ])) self.avg_pool = AdaptiveAvgPool2D(1) - if self.expand: + if self.use_last_conv: self.last_conv = Conv2D( in_channels=make_divisible(NET_CONFIG["stage4"][0] * 2 * scale), - out_channels=self.expand, + out_channels=self.class_expand, kernel_size=1, stride=1, padding=0, @@ -352,18 +309,16 @@ class PPLCNetV2(TheseusLayer): self.dropout = Dropout(p=dropout_prob, mode="downscale_in_infer") self.flatten = nn.Flatten(start_axis=1, stop_axis=-1) - in_features = self.expand if self.expand else NET_CONFIG["stage4"][ - 0] * 2 * scale + in_features = self.class_expand if self.use_last_conv else NET_CONFIG[ + "stage4"][0] * 2 * scale self.fc = Linear(in_features, class_num) def forward(self, x): x = self.stem(x) - x = self.stage1(x) - x = self.stage2(x) - x = self.stage3(x) - x = self.stage4(x) + for stage in self.stages: + x = stage(x) x = self.avg_pool(x) - if self.expand: + if self.use_last_conv: x = self.last_conv(x) x = self.act(x) x = self.dropout(x) @@ -395,6 +350,7 @@ def PPLCNetV2_base(pretrained=False, use_ssld=False, **kwargs): Returns: model: nn.Layer. Specific `PPLCNetV2_base` model depends on args. """ - model = PPLCNetV2(scale=1.0, depths=[2, 2, 6, 2], **kwargs) + model = PPLCNetV2( + scale=1.0, depths=[2, 2, 6, 2], dropout_prob=0.2, **kwargs) _load_pretrained(pretrained, model, MODEL_URLS["PPLCNetV2_base"], use_ssld) return model