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:
parser = argparse.ArgumentParser()
parser.add_argument('--gpu_idx', type=int, default=1)
args = parser.parse_args()
# train model with the learning rate
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.