search

mllam / neural-lam

Neural Weather Prediction for Limited Area Modeling
MIT License
64 stars 24 forks source link

Normalize Data on GPU #25

Open sadamov opened 1 month ago

sadamov commented 1 month ago

Motivation Data normalization can be done on the fly on GPU for each batch. It's faster on GPU than CPU and cleans up the dataset init method.

Implementation Could very nicely use https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#on-after-batch-transfer to normalize once data is on GPU. Makes sure that you never forget about it (all batches on GPU are normalized).

The stats could be provided by a yaml_object handler that can be accessed on the model's init

leifdenby commented 1 month ago

sounds cool @sadamov, are you thinking this for v0.3.0 or a later release? :)

sadamov commented 1 month ago

This feature is ready in https://github.com/mllam/neural-lam/pull/39 I don't have a strong opinion about the version it should be published in. :)