diff --git a/ppcls/arch/__init__.py b/ppcls/arch/__init__.py index da21e101a..3f94501fb 100644 --- a/ppcls/arch/__init__.py +++ b/ppcls/arch/__init__.py @@ -35,8 +35,11 @@ __all__ = ["build_model", "RecModel", "DistillationModel", "AttentionModel"] def build_model(config): arch_config = copy.deepcopy(config["Arch"]) model_type = arch_config.pop("name") + use_sync_bn = arch_config.pop("use_sync_bn", False) mod = importlib.import_module(__name__) arch = getattr(mod, model_type)(**arch_config) + if use_sync_bn: + arch = nn.SyncBatchNorm.convert_sync_batchnorm(arch) if isinstance(arch, TheseusLayer): prune_model(config, arch) quantize_model(config, arch) diff --git a/ppcls/arch/backbone/legendary_models/pp_lcnet.py b/ppcls/arch/backbone/legendary_models/pp_lcnet.py index 3e3c9941f..88dec2041 100644 --- a/ppcls/arch/backbone/legendary_models/pp_lcnet.py +++ b/ppcls/arch/backbone/legendary_models/pp_lcnet.py @@ -321,7 +321,7 @@ def _load_pretrained(pretrained, model, model_url, use_ssld): ) -def PPLCNet_x0_25(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x0_25(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x0_25 Args: @@ -337,7 +337,7 @@ def PPLCNet_x0_25(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs) return model -def PPLCNet_x0_35(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x0_35(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x0_35 Args: @@ -353,7 +353,7 @@ def PPLCNet_x0_35(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs) return model -def PPLCNet_x0_5(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x0_5(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x0_5 Args: @@ -369,7 +369,7 @@ def PPLCNet_x0_5(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): return model -def PPLCNet_x0_75(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x0_75(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x0_75 Args: @@ -385,7 +385,7 @@ def PPLCNet_x0_75(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs) return model -def PPLCNet_x1_0(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x1_0(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x1_0 Args: @@ -401,7 +401,7 @@ def PPLCNet_x1_0(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): return model -def PPLCNet_x1_5(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x1_5(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x1_5 Args: @@ -417,7 +417,7 @@ def PPLCNet_x1_5(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): return model -def PPLCNet_x2_0(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x2_0(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x2_0 Args: @@ -433,7 +433,7 @@ def PPLCNet_x2_0(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): return model -def PPLCNet_x2_5(pretrained=False, use_ssld=False, use_sync_bn=False, **kwargs): +def PPLCNet_x2_5(pretrained=False, use_ssld=False, **kwargs): """ PPLCNet_x2_5 Args: