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

runzhliu

发布于 2020-08-06 10:43:29

3.8K0

文章被收录于专栏：容器计算容器计算

1 Overview

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

2 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_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 这些库是跑步起来的。

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

在 Kubernetes 里运行一个。

apiVersion: v1
kind: Pod
metadata:
  name: tensorflow-gpu
  labels:
    app: tensorflow-gpu
spec:
  containers:
  - name: tensorflow-gpu
    image: tensorflow-gpu-test