Tensorflow 测试一段能运行在 GPU 的代码
147 浏览量
2023-04-13
18:44:10
上传
评论
收藏 15KB DOCX 举报
Example
例子来源于 gihub 上的一段 code，test_single_gpu.py，核心代码很简单，就是在第一块
GPU 上做一个矩阵的运算。
'''
Single GPU computing
'''
with tf.device('/gpu:0'):
a = tf.placeholder(tf.float32, [10000, 10000])
b = tf.placeholder(tf.float32, [10000, 10000])
# Compute A^n and B^n and store results in c1
c1.append(matpow(a, n))
c1.append(matpow(b, n))
with tf.device('/cpu:0'):
sum = tf.add_n(c1) #Addition of all elements in c1, i.e. A^n + B^n
t1_1 = datetime.datetime.now()
with
tf.Session(config=tf.ConfigProto(log_device_placement=log_device_place
ment)) as sess:
# Run the op.
sess.run(sum, {a:A, b:B})
t2_1 = datetime.datetime.now()
复制
Config 里的 log_device_placement 是用来将设备上对 Tensor 的各种操作打印
出来。
Enabling device placement logging causes any Tensor allocations or operations to be
printed.
然后将这份代码放到 Tensorflow 的官方镜像里，docker build 一下，记得要选
GPU 的镜像，否则没有 CUDA 这些库是跑步起来的。
FROM tensorflow/tensorflow:1.14.0-gpu-py3
COPY test_single_gpu.py /
CMD ["python", "/test_single_gpu.py"]
资源评论
