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

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 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()


# train model with the learning rate

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

method: grid
    value: 8
    value: 0.0001
  - ${env}
  - python
  - ${program}
  - "--gpu_idx=1"

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

Thank you very much! I will try it out.

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

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?

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.

1 Like

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.

Hi Thanos, thank you very much for the help all the way! I appreciate it. Setting CUDA_VISIBLE_DEVICES is a viable way for me.

Hi @liu97 perfect, glad to hear this! I will now close this ticket, and please feel free to contact us again if you have any other questions.

This topic was automatically closed 60 days after the last reply. New replies are no longer allowed.