溫馨提示×

tensorflow中怎么指定gpu運行

小億
245
2024-03-27 12:37:31

在TensorFlow中，可以通過以下方式指定GPU運行：

  1. 在代碼中設(shè)置環(huán)境變量：
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 指定使用第一個GPU
  1. 在建立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)
  1. 使用with tf.device()語句指定某段代碼在指定的GPU上運行：
import tensorflow as tf

with tf.device('/device:GPU:0'):
    # 在第一個GPU上運行的代碼
  1. 使用tf.distribute.Strategy指定在多個GPU上運行：
import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 在多個GPU上運行的代碼

注意：在使用GPU時，需要安裝對應(yīng)版本的CUDA和cuDNN，并且確保TensorFlow安裝時支持GPU。

0 踩