在Docker Compose上使用GPU运行TensorFlow

薛才杰 2018-04-19 10:23:00

容器化和AI是目前开发的大趋势。理想情况下，容器化可以将环境无缝迁移，将配置环境的成本无限降低：但是在容器中配置CUDA并运行TensorFlow一向是个难题。对于初学者以及没有深度学习工作站的用户，AWS和Azure推出了带独立显卡的云服务：但是按需实例价格不便宜，竞价式实例价格公道然而虚拟机不能重启，导致不能按需挂载硬盘并保存工作状态，用户需要编写大量代码时刻对虚拟机进行监控并对结果进行异地保存。Deep Systems对这个问题进行了研究并给出了很好的建议。


  1. 用户首先需要配置 CUDA：Deep Systems推荐使用runtime 安装，以免在升级后污染依赖环境。

  2. 安装 Docker：一行命令就可以解决这个问题。curl -sSL https://get.docker.com/ | sh

  3. 安装 Nvidia Docker：Nvidia对Docker有良好的支持，可以方便地配置GPU穿透，而且不必再担心驱动版本不匹配的问题。 

wget -P /tmp https://github.com/NVIDIA/nvidia-docker/releases/download/v1.0.1/nvidia-docker_1.0.1-1_amd64.deb
sudo dpkg -i /tmp/nvidia-docker*.deb

安装后的调用方法是：nvidia-docker run --rm nvidia/cuda nvidia-smi


安装 Docker Compose：Docker Compose可以提供很多方便的功能，即使只有一个容器也推荐使用。 


curl -L https://github.com/docker/compose/releases/download/1.15.0/docker-compose-`uname -s`-`uname -m` > /usr/local/bin/docker-compose
chmod +x /usr/local/bin/docker-compose

安装 Nvidia Docker Compose：这个小脚本可以将Nvidia Docker和Docker Compose连接起来。


使用pip install nvidia-docker-compose 安装后，即可使用nvidia-docker-compose 命令直接调用。


加入别名：nvidia-docker-compose 比较难敲，所以配置下列别名： 


alias doc='nvidia-docker-compose'
alias docl='doc logs -f --tail=100'


最后，在docker-compose.yml 写入下列内容： 


version: '3'
services:
  tf:
    image: gcr.io/tensorflow/tensorflow:latest-gpu
    ports:
      - 8888:8888
    volumes:
      - .:/notebooks


大功告成！doc up 会启动容器，自带TensorFlow Jupiter。


Deep Systems已经将这套系统投入生产。


对于AWS的竞价实例，由于竞价实例不能重启，而且不能保存状态，情况略微复杂一些：有几种聪明的解决方式。


  • 建立一个数据卷，启动时手工挂载。优点是稳定；缺点是每次启动时都必须人工操作，而且数据卷外的所有操作都会丢失。

  • 使用Docker，在启动后挂载一个数据卷，或者按本文方法在启动后建立Docker，直接在里面进行操作。优点是易于理解、工具成熟，缺点是有点麻烦，而且除非按本文方法，否则没有真正解决数据集外操作全部丢失的问题。

  • 使用ec2-spotter，准备一个配置好的数据卷，在启动时直接将系统卷替换：缺点是对于每个区域都需要人工修改所需镜像。

© docker中文社区站