How do I use TensorFlow GPU?

前端 未结 7 1898
予麋鹿
予麋鹿 2020-11-28 06:00

How do I use TensorFlow GPU version instead of CPU version in Python 3.6 x64?

import tensorflow as tf

Pytho

7条回答
  •  甜味超标
    2020-11-28 06:42

    Uninstall tensorflow and install only tensorflow-gpu; this should be sufficient. By default, this should run on the GPU and not the CPU. However, further you can do the following to specify which GPU you want it to run on.

    If you have an nvidia GPU, find out your GPU id using the command nvidia-smi on the terminal. After that, add these lines in your script:

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = #GPU_ID from earlier
    
    config = tf.ConfigProto()
    sess = tf.Session(config=config)
    

    For the functions where you wish to use GPUs, write something like the following:

    with tf.device(tf.DeviceSpec(device_type="GPU", device_index=gpu_id)):
    

提交回复
热议问题