项目作者: afritzler

项目描述 :
Run Tensorflow and Keras with GPU support on Kubernetes
高级语言:
项目地址: git://github.com/afritzler/docker-tensorflow-keras-gpu.git
创建时间: 2017-03-20T09:34:54Z
项目社区:https://github.com/afritzler/docker-tensorflow-keras-gpu

开源协议:

下载


docker-tensorflow-keras-gpu

Base image with Tensorflow and Keras with GPU support. The purpose of this project is to run a single Pod on Kubernetes on a GPU backed node.

Run it on Kubernetes

First, we need to label the gpu instance (if not already done). If you choose to use a different labeling, you need to adjust the nodeSelector part in the example-pod.yaml file.

  1. kubectl label node <NODE_NAME> gpu="true"

Tricky part ahead:

In order to use the GPU inside your docker container, you need to pass the location of the NVidia driver on the host into your container. Since the GPU kernel driver on the host has to match the nvidia-driver inside the contianer, we want to decouple that. Adjust the path in the example-pod.yaml file in case your nvidia-driver location is different.

  1. volumes:
  2. - name: nvidia-driver
  3. hostPath:
  4. path: /var/lib/nvidia-docker/volumes/nvidia_driver/latest

Now run a simple deployment

  1. kubectl create -f example-pod.yaml

Attach to the running container

  1. kubectl exec -it tensorflow-keras-gpu -- /bin/bash

Lets do some training

Inside the docker container I placed the Keras examples from github. To run a simple training example on the IMDB dataset

  1. cd /keras/example
  2. python imdb_cnn.py

If your configuration and driver mapping was done correctly, you should see something like that before the training starts.

  1. Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX TITAN X, pci bus id: 0000:09:00.0)

References