目录

Tensorflow测试一段能运行在GPU的代码

概述

官方文档「又长又臭」,我只是想在 Kubernetes 集群里,运行一个能跑在 GPU 显卡的程序而已,文档太多,看的眼花缭乱,本文就讲一个简单的例子。

Example

例子来源于 github 上的一段 code,test_single_gpu.py,核心代码很简单,就是在第一块 GPU 上做一个矩阵的运算。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
'''
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_placement)) 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 这些库是跑步起来的。

1
2
3
FROM tensorflow/tensorflow:1.14.0-gpu-py3
COPY test_single_gpu.py /
CMD ["python", "/test_single_gpu.py"]

在 Kubernetes 里运行一个。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
apiVersion: v1
kind: Pod
metadata:
  name: tensorflow-gpu
  labels:
    app: tensorflow-gpu
spec:
  containers:
  - name: tensorflow-gpu
    image: tensorflow-gpu-test

总结

测试一段 GPU 的代码,将代码放到合适版本的 TensorFlow 官方的 GPU 镜像,然后通过 Kubernetes 运行起来即可,当然其中需要配置好的 nvidia-docker 之类的环境,本文就不多赘述了。

警告
本文最后更新于 2017年2月1日,文中内容可能已过时,请谨慎参考。