NVIDIA
🚀[FEA]: Allow passing models as to EDMPrecond
Is this a new feature, an improvement, or a change to existing functionality?
Improvement
How would you describe the priority of this feature request
Medium
Please provide a clear description of problem you would like to solve.
Currently models are passed as names (strings) and splatted arguments. This makes it confusing to determine which arguments get passed to the model and which go to the preconditioner. It would be nice to also allow just passing an initialized model itself. The downside is that the user would be required to determine the correct input and output channels, but I don't think this is too onerous (and it could also be checked when the preconditioner is initialized.
Describe any alternatives you have considered
No response
Thanks for reporting @simonbyrne. We will consider your proposed enhancement in the multiple refactorings of the EDMPrecond and the underlying models that are underway.
It would be nice to also allow just passing an initialized model itself. The downside is that the user would be required to determine the correct input and output channels, but I don't think this is too onerous
@simonbyrne I have been looking at implementing this enhancement, and another more serious downside, is that passing directly an initialized model breaks our checkpoint loading/saving, because our checkpoints save all args that are passed to the __init__ constructor of a Module, such that it is able to re-initialize it at load time. But an initialized model cannot be saved in the checkpoint itself because it is not json serializable. So, this option is not going to work.
The only alternative I see to be able to pass any model to EDMPrecond, would be a purely functional interface, where edm_precond is a function instead of a physicsnemo Module. This is a little less convenient, but at least it can be used to wrap any model architecture, and the user would save the model architecture itself instead of the EDM wrapper.
