tensorflow 指定GPU第几片

最新推荐文章于 2023-09-26 17:40:18 发布

算法恩仇录

最新推荐文章于 2023-09-26 17:40:18 发布

阅读量325

点赞数

CC 4.0 BY-SA版权

分类专栏： TensorFlow

本文链接：https://blog.youkuaiyun.com/u010751974/article/details/89738319

TensorFlow 专栏收录该内容

3 篇文章

订阅专栏

博客介绍了TensorFlow中GPU的使用方法。首先展示了指定单块GPU运行的代码，通过设置环境变量指定第0块GPU。接着给出多GPU使用示例，利用递归思想求解矩阵幂次方，将不同计算任务分配到不同GPU（如3、4、7号），最后使用tf.add_n函数求和并输出结果。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1，指定那片GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" #指定在第0块GPU上跑

2，多GPU使用

import numpy as np
import tensorflow as tf
n = 10
A = np.random.rand(2, 2).astype('float32')
B = np.random.rand(2, 2).astype('float32')
c1 = []
# 递归的思想求解幂次方
def matpow(M,n):
if n < 1:
return M
else:
return tf.matmul(M,matpow(M,n-1))
#分配第一个GPU：3
with tf.device('/gpu:3'):
a = tf.placeholder(tf.float32,[2, 2])
c1.append(matpow(a, n))
#分配第二GPU：4
with tf.device('/gpu:4'):
b = tf.placeholder(tf.float32, [2, 2])
c1.append(matpow(b, n))
# tf.add_n([p1, p2, p3....])函数是实现一个列表的元素的相加
with tf.device('/gpu:7'):
sum = tf.add_n(c1)