Docker中使用深度学习框架并支持GPU加速
Docker中使用深度学习框架并支持GPU加速
启动一个支持gpu的容器
docker run --runtime=nvidia --restart=always --name tensorflow -dit -v `pwd`:/app -w /app nvidia/cuda:9.0-cudnn7-runtime-ubuntu16.04
进入容器
docker exec -it tensorflow bash
设置源
Jermine@ubuntu:~$ cat > /etc/apt/sources.list
deb http://mirrors.aliyun.com/ubuntu/ xenial main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ xenial-security main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ xenial-updates main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ xenial-backports main restricted universe multiverse
##测试版源
deb http://mirrors.aliyun.com/ubuntu/ xenial-proposed main restricted universe multiverse
更新源
sudo apt-get update
安装相关依赖
# 导入环境变量
TENSORFLOW_VERSION=1.7.0
apt-get update -y && apt-get install -y --no-install-recommends python3 python3-pip protobuf-compiler;\
pip3 install --upgrade pip ;\
python3 -V && pip3 -V ;\
pip3 --no-cache-dir install setuptools ;\
pip3 --no-cache-dir install \
https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-${TENSORFLOW_VERSION}-cp35-cp35m-linux_x86_64.whl ;\
apt-get autoremove && apt-get autoclean ;\
rm -rf /var/lib/apt/lists/*
测试程序
import numpy as np
np.random.seed(0)
import tensorflow as tf
import time
N,D = 6000,8000
with tf.device('/cpu:0'):
x = tf.placeholder(tf.float32)
y = tf.placeholder(tf.float32)
z = tf.placeholder(tf.float32)
a = x * y
b = a + z
c = tf.reduce_sum(b)
grad_x, grad_y, grad_z = tf.gradients(c, [x,y,z])
start_time = time.time()
with tf.Session() as sess:
values = {
x: np.random.randn(N, D),
y: np.random.randn(N, D),
z: np.random.randn(N, D),
}
out = sess.run([c, grad_x, grad_y, grad_z],
feed_dict=values)
c_val, grad_x_val, grad_y_val, grad_z_val = out
elapsed = time.time() - start_time
print(time.strftime("%H:%M:%S", time.gmtime(elapsed)))
print("exit 0")
将其存为 test_gpu_for_tensorflow.py ， 使用 python3 test_gpu_for_tensorflow.py 执行结果如下：