CODEWITHSUNDEEP
tensorflowinstallationgpu
Will O.
Will O.

Reputation: 11

Specify GPU to run training on (TF object detection API, model zoo)

I am training object detection on a device with multiple GPU's, and want to run training on gpu 1 (keeping 0 and 2 free) and cannot see an option to do so when starting training. I have looked through train.py and model_main.py and cannot find a line to change in there as well. Any suggestions?

Upvotes: 1

Views: 306

Answers (1)

Tamir Tapuhi
Tamir Tapuhi

Reputation: 416

use CUDA_VISIBLE_DEVICES environment variable. you can do it by adding the following to your command line:

CUDA_VISIBLE_DEVICES=1 python <your-python-script>

it will enable only GPU 1 for TF.

Upvotes: 2

Related Questions