mllam / neural-lam

Neural Weather Prediction for Limited Area Modeling
MIT License
102 stars 37 forks source link

Inference step #45

Open clechartre opened 4 months ago

clechartre commented 4 months ago

Purpose

This is the equivalent of now closed PR #41, which was tailored to the fork MeteoSwiss/neural-lam now adapted to this master repo.

Adding the inference functionality to the model trainer. Allows to feed a prediction dataset onto which we can run the build-in prediction function and generate graphical output as well as numpy arrays.

Code changes:

neural_lam/models/ar_model.py: incorporated the predict_step which overwrites trainer.predict(), and incorporation of the on_predict_epoch_end() function to handle model outputs. Adapted the plotting function to be flexible across test and predict steps. Changed the handling of the global-rank specification to using the decorator rather than having if statements in the function.

neural_lam/models/base_graph_model.py: renamed the previous predict_step() for proper use by the trainer.

neural_lam/weather_dataset.py: added a predict dataset and its respective dataloader.

train_model.py: added the option to launch the inference from the trainer. It now has all multiple possibles branching actions as defined in Pytorch Lightning, train, test, eval, and predict.

Checklist

Before submitting this PR, please make sure:

Review

For the review process follow the guidelines at Checklist