Config refactor: improving and standardizing configs and args

[vict0rsch]

• Remove in-code defaults (all trainer parameters should exist in either the yaml config or the command-line arguments, no more config.get(arg, default))
• Clean-up trainer creation
• Clean up model creation
• WandB projects
• BaseModel class
• End of epoch printing
• Remove submit/summit refs
• Improve epoch time metric
• Set up test for FA preserves pbc distances
• Modify FA after theoretical reflexion
• Eval loop in test_all.py

TBD

• max_num_neighbors varies in the code: 32, 40, 50. Which should we choose? Temporarily setting to 40 in base.yamls

[vict0rsch]

Example functioning ✅ run

python main.py --config-yml=configs/is2re/all/forcenet/new_forcenet.yml --mode train --model.phys_embeds=true --model.phys_hidden_channels=32 --model.pg_hidden_channels=32 --model.tag_hidden_channels=64 --optim.max_epochs=10 --wandb_tags=debug --note='developping better-config branch'

[vict0rsch]

Simplified main

trainer: BaseTrainer = registry.get_trainer_class(trainer_config["trainer"])(
            **trainer_config
        )

Simplified Base trainer

@registry.register_trainer("base")
class BaseTrainer(ABC):
    def __init__(self, **kwargs):

        run_dir = kwargs["run_dir"]
        model_name = kwargs["model"].pop("name")

        self.config = {
            **kwargs,
            "model_name": model_name,
            "gpus": distutils.get_world_size() if not kwargs["cpu"] else 0,
            "commit": get_commit_hash(),
            "checkpoint_dir": str(Path(run_dir) / "checkpoints"),
            "results_dir": str(Path(run_dir) / "results"),
            "logs_dir": str(Path(run_dir) / "logs"),
        }
		...

Simplified energy trainer

@registry.register_trainer("energy")
class EnergyTrainer(BaseTrainer):
    """
    Trainer class for the Initial Structure to Relaxed Energy (IS2RE) task.

    .. note::

        Examples of configurations for task, model, dataset and optimizer
        can be found in `configs/ocp_is2re <https://github.com/Open-Catalyst-Project/baselines/tree/master/configs/ocp_is2re/>`_. # noqa: E501
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs, name="is2re")

    def load_task(self):
        logging.info(f"Loading dataset: {self.config['task']['dataset']}")
        self.num_targets = 1