From 4a54d7e690897dd6e6c719fb908fd371a44c2952 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Sun, 24 Jan 2021 22:14:17 +0100 Subject: Many updates, cool stuff on the way. --- src/text_recognizer/networks/util.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) (limited to 'src/text_recognizer/networks/util.py') 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) -- cgit v1.2.3-70-g09d2