1. 程式人生 > >在Docker Compose上使用GPU執行TensorFlow

在Docker Compose上使用GPU執行TensorFlow

容器化和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,準備一個配置好的資料卷,在啟動時直接將系統卷替換:缺點是對於每個區域都需要人工修改所需映象。