❓ [QUESTION] Sweeping hyperparemeters with Weights and Biases

[rschireman]

Hi all,

I'm relatively new to the code here, is there a way to use Weights and Biases to sweep hyperparameters (like the batch size, etc.)? I've been using the following code:

sweep_id = wandb.sweep(sweep_config, project="sweep")
trainer = TrainerWandB(model=model,**dict(minimal_config))
trainer.save()
trainer.set_dataset(dataset)
wandb.agent(sweep_id, trainer.train(), count=5)

and my sweep_config looks like this:

{'method': 'random',
 'metric': {'goal': 'minimize', 'name': 'validation_e'},
 'parameters': {'batch_size': {'distribution': 'q_log_uniform_values',
                               'max': 256,
                               'min': 32,
                               'q': 8},
                'learning_rate': {'distribution': 'uniform',
                                  'max': 0.1,
                                  'min': 0}}}

The first iteration runs fine, but the subsequent runs fail

[Hongyu-yu]

Hi @rschireman,
Sure, with some small modification below on the source code, you can sweep hyperparameters.
In nequip/scripts/train.py line 25, change the default run_name="NequIP", into run_name="rand_"+str(uuid.uuid4()), and insert import uuid in the first line so that the sweep will generate random run_name for the calculations.
You can either change the code in your environment directly or on the source code with pip install ..

Then you can create two config file, one for the sweep hyperparameter and one for the basic config file.

For sweep_config.yaml, here is an example:

project: {your project name}
program: nequip-train
method: random
name: sweep_1
description: sweep over loss coefficients and batch size

# metric:
#   goal: minimize
#   name: validation_f_rmse

command:
  - ${env}
  - ${program}
  - ".../base.yaml" # Absolute path of your basic config yaml file
 
parameters:
  batch_size:
    distribution: int_uniform
    min: 1
    max: 5
  loss_coeffs:
      parameters:
          stress:
              distribution: uniform
              min: 1.0
              max: 10.0
          total_energy:
              value: 
              - 1.
              - PerAtomMSELoss
          forces:
              distribution: uniform
              min: 1
              max: 10

Hope it helps!

[rschireman]

This works, thank you!

[Hongyu-yu]

Glad to hear that :)

[Linux-cpp-lisp]

Thanks very much for posting this info @Hongyu-yu !

An alternative, without modifying the source, would then also be to make your program in for wandb a simple python script that gets the JSON config of the sweep hyperparams, reads the base config, and writes a new YAML file of the merge for the specific run. This script could be re-used for all sweeps.

If that is the only modification needed to make it work, alternatively, we could add a --uuid-run-name flag to nequip-train so this works out of the box. The only unpleasant thing about that is that restarting/finding training runs becomes a pain, since you have to go through each uuid named dir looking for the right values in config.yaml.

[rschireman]

Awesome, I'll see what I can come up with and submit a PR this week

[Linux-cpp-lisp]

Great! Thanks 😄

[Linux-cpp-lisp]

Something useful might be our Config class, i.e. something along the lines of:

from nequip.utils import Config
base_config = Config.from_file("whatever.yaml")
base_config.update(new_params) 
base_config.save("whatever_new.yaml", format="yaml")

[Hongyu-yu]

The only unpleasant thing about that is that restarting/finding training runs becomes a pain, since you have to go through each uuid named dir looking for the right values in config.yaml.

Yes, it may be unconvinient with a random name. One possible solution could be assign run_name with the new_params, for example batch_size:1-stress-3.31-force-2.13-time_tag which may be more straightforward and considered to be included in the script @rschireman . Though if too many hyperparameters are swept, folder name may be too complicated. Also, new_params can be considered to print at first in the log.

[Linux-cpp-lisp]

Mm, I like both of these ideas! (Print in the log and name the run after the hyps--- just have to be careful to correctly sanitize values in a sane way for the filesystem... maybe it's an option but can also use just a simple uuid?)

[mkphuthi]

I'd be interested to see how this was resolved. I originally had a script on <0.5.5 that imported fresh_start from nequip.scripts.train but that seems to be broken on develop. I ideally would want to use a wandb agent to do the sweeping automatically. Does rewriting the config achieve that? I can't seem to justify how that would work.