From d7286199f610a541d092e36694c65f20d62443e9 Mon Sep 17 00:00:00 2001 From: Felipe Bonchristiano Date: Wed, 26 Aug 2026 11:39:02 -0500 Subject: [PATCH] Fix MoleRec fallback predictor never training simple_predictor was created lazily inside forward(), so an optimizer built from model.parameters() beforehand (the standard pattern) never saw its parameters. Create it in __init__ on the no-SMILES fallback path. --- pyhealth/models/molerec.py | 18 +++++++++++------- tests/core/test_molerec.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/pyhealth/models/molerec.py b/pyhealth/models/molerec.py index 559804c8c..b4100d0cc 100644 --- a/pyhealth/models/molerec.py +++ b/pyhealth/models/molerec.py @@ -628,7 +628,14 @@ def __init__( # check if we have valid molecular data (not just dummy data) self.has_valid_smiles = len(self.substructure_smiles) > 1 - + + # Fallback predictor used when there is no valid molecular data. It must + # be created here (not lazily in forward) so its parameters are + # registered before the optimizer is built; otherwise it never trains. + # patient_emb is [condition_emb | procedure_emb], i.e. hidden_dim * 2. + if not self.has_valid_smiles: + self.simple_predictor = torch.nn.Linear(hidden_dim * 2, self.label_size) + self.substructure_graphs = StaticParaDict( **graph_batch_from_smiles(self.substructure_smiles) ) @@ -810,12 +817,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: batch_indices = torch.arange(patient_emb.size(0)).to(self.device) last_patient_emb = patient_emb[batch_indices, last_visit_mask] - # simple linear projection from patient embedding to drug predictions - if not hasattr(self, 'simple_predictor'): - self.simple_predictor = torch.nn.Linear( - patient_emb.size(-1), self.label_size - ).to(self.device) - + # simple linear projection from patient embedding to drug + # predictions (self.simple_predictor is created in __init__ so its + # parameters are registered with the optimizer). logits = self.simple_predictor(last_patient_emb) y_prob = torch.sigmoid(logits) loss = binary_cross_entropy_with_logits(logits, y_true) diff --git a/tests/core/test_molerec.py b/tests/core/test_molerec.py index 0da86cab8..d68c12001 100644 --- a/tests/core/test_molerec.py +++ b/tests/core/test_molerec.py @@ -133,6 +133,35 @@ def test_output_shapes(self): self.assertEqual(ret["y_prob"].shape, (batch_size, num_labels)) self.assertEqual(ret["y_true"].shape, (batch_size, num_labels)) + def test_fallback_predictor_registered_and_trains(self): + """The fallback predictor must be registered before the optimizer. + + Regression test: on the no-valid-SMILES fallback path, simple_predictor + was created lazily inside forward(). An optimizer built from + model.parameters() beforehand (the standard training pattern) therefore + never saw its parameters and never trained it. + """ + self.assertFalse(self.model.has_valid_smiles) + + # simple_predictor exists (and is in named_parameters) before any forward + param_names = [n for n, _ in self.model.named_parameters()] + self.assertTrue(any("simple_predictor" in n for n in param_names)) + + optimizer = torch.optim.Adam(self.model.parameters(), lr=0.1) + before = self.model.simple_predictor.weight.detach().clone() + + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) + data_batch = next(iter(train_loader)) + ret = self.model(**data_batch) + optimizer.zero_grad() + ret["loss"].backward() + optimizer.step() + + self.assertFalse( + torch.equal(before, self.model.simple_predictor.weight.detach()), + "fallback predictor was not updated by the optimizer step", + ) + if __name__ == "__main__": unittest.main()