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
836 stars 190 forks source link

[CorrDiff]: refactor corrdiff configs to dataclasses #417

Open nbren12 opened 3 months ago

nbren12 commented 3 months ago
          @mnabian Can we make an issue to refactor all these options to a common dataclass, which we can then instantiate with https://hydra.cc/docs/1.0/tutorials/structured_config/minimal_example/. Using a dataclass will allow automatic completion, a single point for documentation, and more generally facilitate validation/static checking.

_Originally posted by @nbren12 in https://github.com/NVIDIA/modulus/pull/401#discussion_r1537923976_

nbren12 commented 3 months ago

@mnabian I could make this refactor if you like. Thoughts?

For example, this code:

    # Parse deterministic sampler options
    sigma_min = getattr(cfg, "sigma_min", None)
    sigma_max = getattr(cfg, "sigma_max", None)
    rho = getattr(cfg, "rho", 7)
    solver = getattr(cfg, "solver", "heun")
    discretization = getattr(cfg, "discretization", "edm")
    schedule = getattr(cfg, "schedule", "linear")
    scaling = getattr(cfg, "scaling", None)
    S_churn = getattr(cfg, "S_churn", 0)
    S_min = getattr(cfg, "S_min", 0)
    S_max = getattr(cfg, "S_max", float("inf"))
    S_noise = getattr(cfg, "S_noise", 1)

would change to

@dataclass
class SamplerConfig:
    sigma_min: Optional[float] = None
    sigma_max: Optional[float] = None
    rho: int = 7
    solver: str = "heun"
    discretization: str = "edm"
    schedule: str = "linear"
    scaling: Optional[float] = None
    S_churn: int = 0
    S_min: int = 0
    S_max: float = float("inf")
    S_noise: int = 1
mnabian commented 3 months ago

Yes I think this will clean up the configs. Please go ahead and let me know if I can help.

DavidLandup0 commented 2 months ago

@nbren12 is this still WIP? :) I'd volunteer take it off your plate if not.

nbren12 commented 2 months ago

Go for it @DavidLandup0 !

DavidLandup0 commented 2 months ago

@nbren12 thanks! I'll open a PR later today or tomorrow :)