在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'
的参数即可。