在TensorFlow中使用GPU非常簡單,只需按照以下步驟進行操作:
import tensorflow as tf
# 創建一個TensorFlow會話
with tf.Session() as sess:
# 設置GPU設備
with tf.device('/gpu:0'):
# 這里可以進行GPU上的計算操作
通過以上步驟,你的TensorFlow程序就會在GPU上運行,加快計算速度。值得注意的是,如果你有多個GPU,你可以選擇在不同的GPU上運行不同的操作,只需要修改'/gpu:0'
的參數即可。