Parse additional arguments to the program when running the wandb agent command from the command line

W&B Help
,
1

Is it possible to parse additional arguments to the program when running the wandb agent command from the command line?

For example, suppose I have a script train.py that takes a --gpu_idx argument to specify the GPU index, and I want to run the script with different GPUs using the WandB agent. Can I pass the --gpu_idx argument as a key-value pair when running the wandb agent command?

wandb agent <ID> --gpu_idx 1

In the training script, I have something like:

import wandb
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--gpu_idx', type=int, default=1)
args = parser.parse_args()

wandb.init()
wandb.config.update(args)

# train model with the learning rate
2

Hi @liu97 thank you for writing in! You could parse arguments to your train.py script, you will need to edit your yaml file such that it contains the command field too as follows:

program:
  train.py
method: grid
parameters:
  batch_size:
    value: 8
  lr:
    value: 0.0001
command:
  - ${env}
  - python
  - ${program}
  - "--gpu_idx=1"

Would this work for you? Please also check our documentation here on the command.

3

Thank you very much! I will try it out.

4

Hi @liu97 just checking in here to see if that resolved the issue, and if you had any other questions? thanks

5

Hi, I didn’t run hyper-parameter search with wandb sweep recently, so did not have the chance to test it. I will post updates here.

Just a question, after creating a sweep agent using the above example yaml file, it will use --gpu_idx=1 to sweep through hyper parameters by default. If I run this agent on other machines, where GPU 0 instead of GPU 1 is available, would I be able to parse this --gpu_idx=0 when running wandb agent sweep_id --gpu_idx=0?

6

Hi @liu97 thanks for the additional context. In that case you won’t be able to do this as you will be using the same yaml for all agents. What you could do though in a multi-gpu environment is to specify the GPU as follows:

CUDA_VISIBLE_DEVICES=0 wandb agent sweep_ID

Would this work for you? Please also check this docs page for more information.

7

Hi @liu97 I wanted to follow up with you and see if you’ve managed to run the wandb agent in multiple GPUs? and if there are any further questions that we could help with.