diff --git a/docs/api/processors/pyhealth.processors.TensorProcessor.rst b/docs/api/processors/pyhealth.processors.TensorProcessor.rst index 07a600bb5..3fd9db3b8 100644 --- a/docs/api/processors/pyhealth.processors.TensorProcessor.rst +++ b/docs/api/processors/pyhealth.processors.TensorProcessor.rst @@ -3,6 +3,23 @@ pyhealth.processors.TensorProcessor Processor for tensor data. +``size()`` returns the width of the feature's last dimension seen in ``fit()`` +(1 for scalars), or ``None`` before fitting. + +Custom tensor processors +------------------------ + +Subclass ``TensorProcessor`` to transform numeric features, for example to select +columns, impute or scale. Models such as ``MLP`` size their input layer from the +already-processed samples, falling back to ``size()``, and never call +``process()`` on a sample a second time, so ``process()`` does not need to be +idempotent. If ``process()`` changes the width, override ``size()`` to return the +output width. See ``examples/custom_tensor_processor.py``. + +In PyHealth 2.0.2 and earlier, ``EmbeddingModel`` called ``process()`` again on +already-processed samples to infer the width, which broke processors that are not +idempotent. + .. autoclass:: pyhealth.processors.TensorProcessor :members: :undoc-members: diff --git a/examples/custom_tensor_processor.py b/examples/custom_tensor_processor.py new file mode 100644 index 000000000..4d7732660 --- /dev/null +++ b/examples/custom_tensor_processor.py @@ -0,0 +1,63 @@ +"""A custom tensor processor that changes the feature width, used with an MLP. + +Subclass ``TensorProcessor`` to transform numeric features, here keeping a +subset of columns. Models size their input layer from the already-processed +samples (or from ``size()``), and never call ``process()`` a second time, so a +processor like this one, which is not idempotent, works with any model. + +Runs in a few seconds on CPU with synthetic data; no download. + +Usage: + python examples/custom_tensor_processor.py +""" + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import MLP +from pyhealth.processors import TensorProcessor + + +class SelectColumns(TensorProcessor): + """Keeps the given columns of a numeric feature vector.""" + + def __init__(self, columns): + super().__init__() + self.columns = list(columns) + + def process(self, value): + return super().process(value)[..., self.columns] + + def size(self): + return len(self.columns) + + +def main(): + torch.manual_seed(0) + samples = [ + { + "patient_id": f"p{i}", + # age, heart rate, temperature, a column we want to drop + "vitals": [40.0 + i, 70.0 + i % 9, 36.5 + (i % 3) / 10, 999.0], + "label": int(i % 3 == 0), + } + for i in range(30) + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"vitals": SelectColumns(columns=[0, 1, 2])}, + output_schema={"label": "binary"}, + dataset_name="custom_tensor_processor", + ) + print("processed sample:", dataset[0]["vitals"]) + + model = MLP(dataset=dataset) + print("input layer:", model.embedding_model.embedding_layers["vitals"]) + + batch = next(iter(get_dataloader(dataset, batch_size=8))) + out = model(**batch) + print("y_prob shape:", tuple(out["y_prob"].shape), "loss:", round(float(out["loss"]), 4)) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/models/embedding.py b/pyhealth/models/embedding.py index 4232b2788..446c94bf4 100644 --- a/pyhealth/models/embedding.py +++ b/pyhealth/models/embedding.py @@ -133,9 +133,26 @@ class EmbeddingModel(BaseModel): Input: (..., size) Output: (..., embedding_dim) - - TensorProcessor: nn.Linear (size inferred from first sample) + - TensorProcessor: nn.Linear, sized from the first already-processed + sample (or the processor's ``size()``); ``process()`` is not re-run - MultiHotProcessor: nn.Linear over multi-hot vector + + Examples: + >>> from pyhealth.datasets import create_sample_dataset + >>> from pyhealth.models import EmbeddingModel + >>> samples = [ + ... {"patient_id": "p0", "x": [1.0, 2.0, 3.0], "label": 0}, + ... {"patient_id": "p1", "x": [4.0, 5.0, 6.0], "label": 1}, + ... ] + >>> dataset = create_sample_dataset( + ... samples=samples, + ... input_schema={"x": "tensor"}, + ... output_schema={"label": "binary"}, + ... ) + >>> model = EmbeddingModel(dataset, embedding_dim=8) + >>> model.embedding_layers["x"] + Linear(in_features=3, out_features=8, bias=True) """ def __init__( @@ -213,16 +230,18 @@ def __init__( ) elif isinstance(processor, TensorProcessor): - # Infer size from first sample - sample_tensor = None + # Dataset samples are already processed: read the width from the + # first one. Never call process() again, since processors that + # select columns, scale or impute are not idempotent. + input_size = None for sample in dataset: if field_name in sample: - sample_tensor = processor.process(sample[field_name]) + value = torch.as_tensor(sample[field_name]) + input_size = value.shape[-1] if value.dim() > 0 else 1 break - if sample_tensor is not None: - input_size = ( - sample_tensor.shape[-1] if sample_tensor.dim() > 0 else 1 - ) + if input_size is None: + input_size = processor.size() + if input_size is not None: self.embedding_layers[field_name] = nn.Linear( in_features=input_size, out_features=embedding_dim ) diff --git a/pyhealth/processors/tensor_processor.py b/pyhealth/processors/tensor_processor.py index b1270051d..0e4ab47f0 100644 --- a/pyhealth/processors/tensor_processor.py +++ b/pyhealth/processors/tensor_processor.py @@ -19,6 +19,15 @@ class TensorProcessor(FeatureProcessor): Output: - torch.Tensor with appropriate shape and dtype + + Examples: + >>> from pyhealth.processors import TensorProcessor + >>> processor = TensorProcessor() + >>> processor.fit([{"x": [1.0, 2.0, 3.0]}], "x") + >>> processor.size() + 3 + >>> processor.process([4.0, 5.0, 6.0]) + tensor([4., 5., 6.]) """ def __init__( @@ -37,10 +46,11 @@ def __init__( """ self.dtype = dtype self._n_dim = None + self._feature_dim = None self._spatial_dims = spatial_dims def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: - """Infer n_dim from the first valid sample. + """Infer n_dim and the feature width from the first valid sample. Args: samples: Iterable of sample dictionaries. @@ -55,6 +65,7 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: else torch.tensor(value, dtype=self.dtype) ) self._n_dim = tensor.dim() + self._feature_dim = tensor.shape[-1] if tensor.dim() > 0 else 1 break def process(self, value: Any) -> torch.Tensor: @@ -74,14 +85,15 @@ def process(self, value: Any) -> torch.Tensor: return value.detach().clone().to(dtype=self.dtype) return torch.tensor(value, dtype=self.dtype) - def size(self) -> None: + def size(self) -> int | None: """ Get the feature size of the processor. Returns: - None: Size is not predetermined for tensor processor + The width of the last dimension (1 for scalars) seen in ``fit()``, + or None before fitting. """ - return None + return getattr(self, "_feature_dim", None) def is_token(self) -> bool: """Whether the output tensor represents discrete token indices, inferred from dtype. diff --git a/tests/core/test_embedding_tensor_width.py b/tests/core/test_embedding_tensor_width.py new file mode 100644 index 000000000..2cbd1d198 --- /dev/null +++ b/tests/core/test_embedding_tensor_width.py @@ -0,0 +1,80 @@ +"""EmbeddingModel must size tensor features without re-processing samples. + +Samples read from a SampleDataset are already processed. Calling +``processor.process()`` on them again is wrong for any processor that is not +idempotent (one that selects columns, scales or imputes). +""" + +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset +from pyhealth.models import MLP +from pyhealth.processors import TensorProcessor + + +def _samples(): + return [ + {"patient_id": f"p{i}", "x": [float(i), 1.0, 2.0, 3.0], "label": i % 2} + for i in range(6) + ] + + +class FirstTwoColumns(TensorProcessor): + """Keeps the first two columns: output width 2, not idempotent.""" + + def process(self, value): + tensor = super().process(value) + if tensor.shape[-1] != 4: + raise ValueError("process() called on an already-processed sample") + return tensor[..., :2] + + def size(self): + return 2 + + +class TestTensorProcessorSize(unittest.TestCase): + def test_fit_records_feature_width(self): + p = TensorProcessor() + p.fit(_samples(), "x") + self.assertEqual(p.size(), 4) + + def test_size_is_none_before_fit(self): + self.assertIsNone(TensorProcessor().size()) + + def test_scalar_feature_has_width_one(self): + p = TensorProcessor() + p.fit([{"x": 3.0}], "x") + self.assertEqual(p.size(), 1) + + +class TestEmbeddingTensorWidth(unittest.TestCase): + def test_builds_without_reprocessing_samples(self): + dataset = create_sample_dataset( + samples=_samples(), + input_schema={"x": FirstTwoColumns()}, + output_schema={"label": "binary"}, + dataset_name="test_embedding_tensor_width", + ) + self.assertEqual(tuple(dataset[0]["x"].shape), (2,)) + model = MLP(dataset=dataset) # raised ValueError before the fix + layer = model.embedding_model.embedding_layers["x"] + self.assertEqual(layer.in_features, 2) + batch = {"x": torch.stack([dataset[i]["x"] for i in range(3)])} + out = model.embedding_model(batch) + self.assertEqual(out["x"].shape[0], 3) + + def test_plain_tensor_feature_width(self): + dataset = create_sample_dataset( + samples=_samples(), + input_schema={"x": "tensor"}, + output_schema={"label": "binary"}, + dataset_name="test_embedding_tensor_width_plain", + ) + model = MLP(dataset=dataset) + self.assertEqual(model.embedding_model.embedding_layers["x"].in_features, 4) + + +if __name__ == "__main__": + unittest.main()