diff --git a/pyhealth/models/cnn.py b/pyhealth/models/cnn.py index 3d5fe3ce0..29437a722 100644 --- a/pyhealth/models/cnn.py +++ b/pyhealth/models/cnn.py @@ -329,6 +329,10 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: x = x.to(self.device) spatial_dim = self.feature_conv_dims[feature_key] + # Treat spatial_dim==1 features (which embed to [batch, embedding_dim] + # with no sequence axis) as a length-1 sequence so the 1D CNN can run on them. + if spatial_dim == 1 and x.dim() == 2: + x = x.unsqueeze(1) expected_dims = {1: 3, 2: 4, 3: 5}[spatial_dim] if x.dim() != expected_dims: raise ValueError( diff --git a/tests/core/test_cnn.py b/tests/core/test_cnn.py index 4a53e5444..e05b18f7a 100644 --- a/tests/core/test_cnn.py +++ b/tests/core/test_cnn.py @@ -239,6 +239,54 @@ def test_model_with_mixed_inputs(self): self.assertEqual(ret["logit"].shape[0], 2) self.assertEqual(ret["loss"].dim(), 0) + def test_model_with_multihot_and_1d_tensor_inputs(self): + """Test CNN model with non-sequence spatial_dim=1 inputs. + + MultiHotProcessor and a TensorProcessor whose per-sample value is a 1D + vector both embed to [batch, embedding_dim] with no sequence axis. These + are documented as supported input types and must not crash. + """ + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "demographics": ["asian", "non_hispanic"], + "vitals": [1.0, 2.5, 3.0], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-1", + "demographics": ["white"], + "vitals": [0.5, 1.0, 2.0], + "label": 0, + }, + ] + + input_schema = {"demographics": "multi_hot", "vitals": "tensor"} + output_schema = {"label": "binary"} + + dataset = create_sample_dataset( + samples=samples, + input_schema=input_schema, + output_schema=output_schema, + dataset_name="test_multihot", + ) + + model = CNN(dataset=dataset) + self.assertEqual(model.feature_conv_dims["demographics"], 1) + self.assertEqual(model.feature_conv_dims["vitals"], 1) + + train_loader = get_dataloader(dataset, batch_size=2, shuffle=False) + data_batch = next(iter(train_loader)) + + ret = model(**data_batch) + ret["loss"].backward() + + self.assertEqual(ret["y_prob"].shape[0], 2) + self.assertEqual(ret["logit"].shape[0], 2) + self.assertEqual(ret["loss"].dim(), 0) + if __name__ == "__main__": unittest.main()