diff --git a/deepethogram/losses.py b/deepethogram/losses.py index 83a192c..bcdb2a0 100644 --- a/deepethogram/losses.py +++ b/deepethogram/losses.py @@ -186,7 +186,7 @@ def get_regularization_loss(cfg: DictConfig, model): pretrained_dir = cfg.project.pretrained_path assert os.path.isdir(pretrained_dir) weights = projects.get_weights_from_model_path(pretrained_dir) - pretrained_file = weights[cfg.run.model][cfg[cfg.run.model].arch] + pretrained_file = weights.get(cfg.run.model, {}).get(cfg[cfg.run.model].arch, []) if len(pretrained_file) == 0: log.warning("No pretrained file found. Regularization: L2. alpha={}".format(cfg.train.regularization.beta)) diff --git a/tests/test_losses.py b/tests/test_losses.py new file mode 100644 index 0000000..40129d4 --- /dev/null +++ b/tests/test_losses.py @@ -0,0 +1,32 @@ +import torch +from omegaconf import OmegaConf + +from deepethogram import losses + + +def test_l2_sp_missing_pretrained_architecture_falls_back_to_l2(tmp_path, monkeypatch): + cfg = OmegaConf.create( + { + "project": {"pretrained_path": str(tmp_path)}, + "run": {"model": "sequence"}, + "sequence": {"arch": "missing_arch"}, + "train": { + "regularization": { + "style": "l2_sp", + "alpha": 0.1, + "beta": 0.2, + } + }, + } + ) + model = torch.nn.Linear(2, 1) + + def no_sequence_weights(_): + return {"sequence": {}} + + monkeypatch.setattr(losses.projects, "get_weights_from_model_path", no_sequence_weights) + + regularization = losses.get_regularization_loss(cfg, model) + + assert isinstance(regularization, losses.L2) + assert regularization.alpha == cfg.train.regularization.beta