在TensorFlow中,可以通过以下方式指定GPU运行:
- 在代码中设置环境变量:
import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 指定使用第一个GPU
- 在建立session时指定使用的GPU:
import tensorflow as tf gpu_options = tf.GPUOptions(visible_device_list="0") # 指定使用第一个GPU config = tf.ConfigProto(gpu_options=gpu_options) sess = tf.Session(config=config)
- 使用with tf.device()语句指定某段代码在指定的GPU上运行:
import tensorflow as tf with tf.device('/device:GPU:0'): # 在第一个GPU上运行的代码
- 使用tf.distribute.Strategy指定在多个GPU上运行:
import tensorflow as tf strategy = tf.distribute.MirroredStrategy() with strategy.scope(): # 在多个GPU上运行的代码
注意:在使用GPU时,需要安装对应版本的CUDA和cuDNN,并且确保TensorFlow安装时支持GPU。