ACEsuit / mace

MACE - Fast and accurate machine learning interatomic potentials with higher order equivariant message passing.
Other
412 stars 155 forks source link

adding cmd argument for specifying wandb directory #433

Closed nqhq-lou closed 1 week ago

nqhq-lou commented 1 month ago

Background: In our current setup, we configure directories for logs, models, checkpoints, and results. They helped a lot with keeping everything ordered. And it is super that we integrated wandb for logging trainings. For me, wandb is integrated into my workflow to manage the training hyperparameters and output values. Typically, these are stored as JSON files (wandb-[metadata|summary].json) under the directory ${workrootdir}/wandb/${wandb_subdir}, where ./wandb/ is the default value of the keyword in wandb.init(), and ${wandb_subdir} is generated on the fly, e.g. run-20240531_154023-9cce39hr.

Problem Statement: The wandb logging directories are stored under ${workrootdir}/wandb/${wandb_subdir}. This setup does not allow for easily identifiable directory names and relate them to mace training names, leading to difficulties in accurately locating the correct wandb directories and making personalized postprocess difficult.

Proposed Solution: I propose to introduce a new command-line argument --wandb_dir that users can specify when running mace/cli/run_train.py. This argument corresponds directly to the dir parameter in the wandb.init() function. By implementing this feature, users will have the flexibility to define custom and more descriptive directory paths for wandb logs, enhancing data organization and accessibility.

ilyes319 commented 3 weeks ago

Hey @nqhq-lou, happy to merge that if you make a PR!

ilyes319 commented 1 week ago

I merged, thank you @nqhq-lou