概述
官方文档「又长又臭」，我只是想在 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.
hugo.oral.wandb.启动流程.service.gorilla.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日，文中内容可能已过时，请谨慎参考。