TensorFlow How to specify the GPU for training when training a model

When using TensorFlow to train deep learning models, assuming that we did not specify a GPU to train before training, the default is to use the No. 0 GPU to train our model, and the other GPU's will be shown to be occupied. Sometimes we prefer to train our models by specifying a piece or a few gpus ourselves, rather than using this default method. The next step is to introduce two simple methods.

The number of GPUs we have available is as follows:

1. The first is to specify the GPU to be used during training by using the Tf.device () function.

Suppose we want to train the model with our 2nd GPU, which can be specified by the following code:

Tf.device ('/gpu:2 ')

The experimental results are as follows:

We can see that although the 2nd GPU is specified for training, several other GPUs are still occupied, while the model is only trained on the 2nd GPU.

2. Specify by Cuda_visible_devices.

Also using the 2nd GPU to train the model, we can add in our Python code:

Import os
os.environ[' cuda_visible_devices ']= ' 2 '

The experimental results are as follows:

We can see that when we train the model, we use the 2nd GPU, and a few other GPUs are not accounted for, which is equivalent to masking the GPU except for the 2nd block when we run the program, only the 2nd GPU is visible to the currently running program. Similarly, if you want to specify 2, 3 GPU to train, then the above code can be changed to:

os.environ[' cuda_visible_devices ']= ' 2,3 '

Similarly, if you have more GPUs to specify, you can add them to the code above.

This method can also be specified before running a python program, such as:

cuda_visible_devices=2 python train.py

The functionality implemented is the same as the code above.

Reference Links:

Click to open link

Click to open link