[samehaisaa]

Contributor: Sameh Aissa

Contribution Type: Model update (2.0 port) and new task

Description:

Part of #1201. On master VAE can't be created. VAE(dataset) raises TypeError: VAE.__init__() missing 5 required positional arguments and the old 1.x call raises TypeError: BaseModel.__init__() got an unexpected keyword argument 'feature_keys'. The layers were also only built for 32, 64 and 128 px images.

I followed jhnwu3's review on #627. VAE now takes a SampleDataset and reads the image feature, its channels and its size from the ImageProcessor, so there's no input_type or size argument anymore. Anything other than a single image feature raises a clear error. It reports regression mode with the input image as the target, so it trains with the standard Trainer and regression metrics like mse and mae score the reconstruction. Images in [0, 1] keep the BCE loss and normalized images use MSE since they can go outside that range. generate(num_samples) decodes samples from the prior.

To use it on real data there's a new COVID19CXRImageGeneration task with one unlabeled 128 x 128 grayscale sample per X-ray, and the CXR example uses it now.

It's unconditional for now and conditioning on the disease label can be a follow up. I kept it in models/ so from pyhealth.models import VAE doesn't change, happy to move it to models/generators/ if you prefer.

Files to Review:

• pyhealth/models/vae.py the port
• pyhealth/tasks/covid19_cxr_generation.py the new task, exported from pyhealth/tasks/__init__.py
• tests/core/test_vae.py 16 tests on small synthetic PNGs, including a fake copy of the COVID-19 dataset layout
• examples/cxr/chextXray_image_generation_VAE.py the example on the 2.0 API
• docs/api/models/pyhealth.models.VAE.rst and the new task page in docs/api/tasks/

Testing: Windows, Python 3.12, torch 2.15 CPU

• pytest tests/core/test_vae.py 16 passed, coverage 97% on vae.py and 100% on the task
• python tools/check_pr_rules.py --base upstream/master --head HEAD passes
• the example runs end to end on a fake copy of the dataset folder

[Logiquo]

/review

[github-actions]

Test Coverage

Total: 73.53% (19690 / 26778 lines)

Modified file	Coverage	Lines
pyhealth/models/vae.py	96.77%	90 / 93
pyhealth/tasks/__init__.py	89.47%	34 / 38
pyhealth/tasks/covid19_cxr_generation.py	100.00%	9 / 9

Workflow run · 0387f19

[github-actions]

_{Reviewed commit 0387f19}

[github-actions]

[low] Public encoder API changed meaning and removed methods without a note

The old VAE exposed encoder(x) returning (mu, std), plus sampling and a static loss_function. The port replaces these with encode(x) returning (mu, logvar), makes encoder an nn.Sequential, and removes sampling and loss_function. Since the old class could not be constructed on 2.0, no working caller breaks, so this is not blocking. Still, the PR description and the VAE docstring should say that encode returns log-variance rather than std and that the old helper methods were removed, so anyone reading the old API is not misled.

[github-actions]

AI Review

suggestion: approve
decision: approve

Thank you for this contribution, Sameh. It brings VAE back to a working state on the 2.0 API. The model reads the image shape from the ImageProcessor and trains with the standard Trainer. It reports reconstruction error through regression metrics and has a generate() method. It also adds the COVID19CXRImageGeneration task, an updated example and docs.

I found no blocking problems, so this looks ready to merge. The diff is self-contained: it doesn't change base classes, processors or the Trainer. The old VAE could not be constructed on 2.0, so no working caller breaks. The tests cover shapes, gradients, eval determinism, odd sizes and RGB, normalized input, ignored labels, error cases, a Trainer fit, a learning check and the task on a fake dataset layout.

Coverage meets the bar. vae.py is at 96.8%, the new task file at 100%, and tasks/init.py at 89.5%. The whole repository is at 73.5%, above the 60% guideline.

Optional suggestions:

• Mention in the PR description or a changelog that encode() now returns (mu, logvar) where the old encoder returned std. Also note that loss_function and sampling were removed.
• Trainer.inference keeps full reconstruction arrays in memory, so very large splits could use a lot of RAM. A short note in the example would help users.

_{Reviewed commit 0387f19}