physicsnemo icon indicating copy to clipboard operation
physicsnemo copied to clipboard

[CorrDiff] Move the training loop into the modulus package

Open nbren12 opened this issue 1 year ago • 1 comments

Specifically, I am proposing this

git mv examples/generative/corrdiff/training/training_loop.py modulus/utils/generative/

This will allow other projects to use the training loop by importing it. We have close duplicates of this file in 3-4 internal research projects, and are encountering the same bugs and adding the same features. A modulus implementation of the training loop would be widely usable across EDM-based diffusion projects.

The training loop is also the most frequently modified part of corrdiff, but it has no unit tests to ensure that we aren't breaking it. I feel some interfaces and extension points added that could make it more testable and general.

The current workaround is to add modulus as a submodule or otherwise modify the python path to point to examples/generative/corrdiff/ in an ad-hoc manner. In practice, most will probably just copy-paste it, which leads to code duplication.

nbren12 avatar Apr 05 '24 01:04 nbren12

@ram-cherukuri could you please review this request?

mnabian avatar Apr 05 '24 15:04 mnabian