How to select GPU in code?

princeton-vl / pose-ae-train

Training code for "Associative Embedding: End-to-End Learning for Joint Detection and Grouping"

BSD 3-Clause "New" or "Revised" License

373 stars 76 forks source link

How to select GPU in code? #13

Closed yw155 closed 6 years ago

yw155 commented 6 years ago

Hi, I would like to ask that how to select specific GPUs for training in scripts. And I see the memory utilization of multiple GPUs are not the same. Does anyone know how to choose and how to balance the memory utilization? Thanks.

yw155 commented 6 years ago

Now I know how to choose GPU, it is like this: CUDA_VISIBLE_DEVICES=2,3 python3 train.py -e test_run_001 Thanks.