diff options
Diffstat (limited to 'src/text_recognizer/networks/util.py')
| -rw-r--r-- | src/text_recognizer/networks/util.py | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/src/text_recognizer/networks/util.py b/src/text_recognizer/networks/util.py index 711a952..131a6b4 100644 --- a/src/text_recognizer/networks/util.py +++ b/src/text_recognizer/networks/util.py @@ -65,13 +65,18 @@ def configure_backbone(backbone: str, backbone_args: Dict) -> Type[nn.Module]: network_args = state_dict["network_args"] weights = state_dict["model_state"] + freeze = False + if "freeze" in backbone_args and backbone_args["freeze"] is True: + backbone_args.pop("freeze") + freeze = True + network_args = backbone_args + # Initializes the network with trained weights. backbone = backbone_(**network_args) backbone.load_state_dict(weights) - if "freeze" in backbone_args and backbone_args["freeze"] is True: + if freeze: for params in backbone.parameters(): params.requires_grad = False - else: backbone_ = getattr(network_module, backbone) backbone = backbone_(**backbone_args) |