NVIDIA / modulus

Open-source deep-learning framework for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods

https://developer.nvidia.com/modulus

Apache License 2.0

937 stars 219 forks source link

ModAFNO for interpolation #547

Closed jleinonen closed 2 months ago

jleinonen commented 3 months ago

Modulus Pull Request

Description

This adds a Modulated AFNO (ModAFNO) layer that adds a modulation input to the AFNO architecture. This can be used to build tunable AFNO models and works, for example, to build a weather interpolation model.

Checklist

[x] I am familiar with the Contributing Guidelines.
[x] New or existing tests cover these changes.
[x] The documentation is up to date with these changes.
[x] The CHANGELOG.md is up to date with these changes.
[ ] An issue is linked to this pull request.

Dependencies

No new dependencies needed

jleinonen commented 3 months ago

/blossom-ci

jleinonen commented 3 months ago

/blossom-ci

jleinonen commented 3 months ago

/blossom-ci

jleinonen commented 3 months ago

/blossom-ci

jleinonen commented 3 months ago

/blossom-ci

NickGeneva commented 3 months ago

Can you add the model to the docs rst file:

https://github.com/NVIDIA/modulus/blob/main/docs/api/modulus.models.rst

jleinonen commented 3 months ago

Can you add the model to the docs rst file:

https://github.com/NVIDIA/modulus/blob/main/docs/api/modulus.models.rst

Added in https://github.com/NVIDIA/modulus/pull/547/commits/5f46f434155c6cf2925eb5e8f0fb4b8cd89f8655, @NickGeneva could you check that this makes sense? Also, shouldn't AdaptiveFourierNeuralOperator in that file be just AFNO?

NickGeneva commented 3 months ago

Over all I think edits look good, confirm that the outputs are consistent with existing checkpoint and I think we can merge

NickGeneva commented 3 months ago

Oh also please add a bullet in the change log about adding this arch.

jleinonen commented 2 months ago

/blossom-ci

jleinonen commented 2 months ago

@NickGeneva I have confirmed that the original checkpoint loads fine into the Modulus version of ModAFNO in this PR, no changes required.

jleinonen commented 2 months ago

@NickGeneva could you also nominate another reviewer, we'll need 2 approvals to merge.

NickGeneva commented 2 months ago

/blossom-ci

jleinonen commented 2 months ago

/blossom-ci