Reveal Max Training Iterations to Guidance Modules

[ayaanzhaque]

Is there an easy way that I can change the code to pass the max training iterations to the guidance module or forward calls?

[bennyguo]

May I ask what's the use case?

[ayaanzhaque]

I'd like to set up some sort of scheduler that depends on the max iterations.

[thuliu-yt16]

The easiest way I come up with is to inherit the guidance and reimplement update_step.

[bennyguo]

@ayaanzhaque You could pass trainer.max_steps to do_update_step and update_step as @thuliu-yt16 suggested, but this requires a lot of rewriting. An easier solution is using OmegaConf resolvers in the configuration file to assign attributes that are related to trainer.max_steps, e.g., n_steps: ${idiv:${trainer.max_steps},2} assigns n_steps to the half of trainer.max_steps. You can define any resolvers you like as in:

https://github.com/threestudio-project/threestudio/blob/3dc16665f3ff6e2d90b5891f611e35436c2407fb/threestudio/utils/config.py#L10-L24