在特定的 gpu 上运行 python 代码以获得较低的 python 版本

Question

I am trying to run a python code on a specific GPU on our server. The server has four GPUs. When I run the code using a virtual environment installed with python 3.8 and tensorflow 2.2, it works correctly on the specific GPU just by adding the below few lines at the first of the script.

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"  # run the code on a specified GPU

Many threads recommend use the above code to run python scripts on a specific GPU such as here and here .

However, When I tried to use the same way to run another python code on another virtual environment (with lower specifications) that was installed with python version 3.6.9 and tensorflow 1.12, it does not run on the GPU but on the CPU.

How can I run python code on a specific GPU in the case of the second virtual environment?

Answer 1

You can use export CUDA_VISIBLE_DEVICES to define which GPUs are visible to the application. For example, if you want GPUs 0 and 2 visible, use export CUDA_VISIBLE_DEVICES=0,2 .