巧用 Docker 快速部署 GPU 環(huán)境

在 Linux 服務(wù)器上使用 GPU 跑深度學(xué)習(xí)的模型很正常不過。如果我們想用 Docker 實(shí)現(xiàn)同樣的需求，就需要做些額外的工作。本質(zhì)上就是我們要在容器里能看到并且使用宿主機(jī)上的顯卡。在這篇文章里我們就介紹一下 Docker 使用 GPU 的環(huán)境搭建。

Nvidia 驅(qū)動(dòng)

某些命令以 Ubuntu 作為示例。首先宿主機(jī)上必現(xiàn)安裝 Nvidia 驅(qū)動(dòng)。

這里推薦從 Nvidia 官網(wǎng)下載腳本安裝，安裝和卸載都比較方便并且適用于任何 Linux 發(fā)行版，包括 CentOS，Ubuntu 等。NVIDIA Telsa GPU 的 Linux 驅(qū)動(dòng)在安裝過程中需要編譯 kernel module，系統(tǒng)需提前安裝 gcc 和編譯 Linux Kernel Module 所依賴的包，例如?kernel-devel-$(uname -r)?等。

安裝 gcc 和 kernel-dev

$?sudo?apt?install?gcc?kernel-dev?-y

安裝 Nvidia 驅(qū)動(dòng)

1. 訪問 https://www.nvidia.com/Download/Find.aspx

2. 選擇對(duì)應(yīng)操作系統(tǒng)和安裝包，并單擊 [SEARCH] 搜尋驅(qū)動(dòng)，選擇要下載的驅(qū)動(dòng)版本

3. 在宿主機(jī)上下載并執(zhí)行對(duì)應(yīng)版本安裝腳本

$?wget?https://www.nvidia.com/content/DriverDownload-March2009/confirmation.php?url=/tesla/450.80.02/NVIDIA-Linux-x86_64-450.80.02.run&lang=us&type=Tesla
$?chmod?+x?NVIDIA-Linux-x86_64-450.80.02.run?&&?./NVIDIA-Linux-x86_64-450.80.02.run

4. 驗(yàn)證

使用?nvidia-smi?命令驗(yàn)證是否安裝成功，如果輸出類似下圖則驅(qū)動(dòng)安裝成功。

CUDA 驅(qū)動(dòng)

CUDA（Compute Unified Device Architecture）是顯卡廠商 NVIDIA 推出的運(yùn)算平臺(tái)。CUDA?是一種由 NVIDIA 推出的通用并行計(jì)算架構(gòu)，該架構(gòu)使 GPU 能夠解決復(fù)雜的計(jì)算問題。它包含了 CUDA 指令集架構(gòu)（ISA）以及 GPU 內(nèi)部的并行計(jì)算引擎。這里安裝的方式和顯卡驅(qū)動(dòng)安裝類似。

1. 訪問官網(wǎng)下載對(duì)應(yīng)版本安裝包，https://developer.nvidia.com/cuda-toolkit-archive
   
   

   

2. 配置環(huán)境變量

$?echo?'export?PATH=/usr/local/cuda/bin:$PATH'?|?sudo?tee?/etc/profile.d/cuda.sh?
$?source?/etc/profile

nvidia-docker2

Docker 的安裝這里就不展開了，具體查看官方文檔非常詳細(xì)。

這里我們就直接介紹安裝 nvidia-docker2.

既然叫 nvidia-docker2 就有 nvidia-docker1 就是它的 1.0 版本目前已經(jīng)廢棄了，所以注意不要裝錯(cuò)。

這里先簡(jiǎn)單說一下 nvidia-docker2 的原理，nvidia-docker2 的依賴由下幾部分組成.

• libnvidia-container
• nvidia-container-toolkit
• nvidia-container-runtime
  
  

  
  

nvidia-container-runtime 是在 runc 基礎(chǔ)上多實(shí)現(xiàn)了 nvidia-container-runime-hook (現(xiàn)在叫 nvidia-container-toolkit)，該 hook 是在容器啟動(dòng)后（Namespace已創(chuàng)建完成），容器自定義命令(Entrypoint)啟動(dòng)前執(zhí)行。當(dāng)檢測(cè)到 NVIDIA_VISIBLE_DEVICES 環(huán)境變量時(shí)，會(huì)調(diào)用 libnvidia-container 掛載 GPU Device 和 CUDA Driver。如果沒有檢測(cè)到 NVIDIA_VISIBLE_DEVICES 就會(huì)執(zhí)行默認(rèn)的 runc。

下面分兩步安裝

1. 設(shè)置 repository 和 GPG key

$?distribution=$(.?/etc/os-release;echo?$ID$VERSION_ID)
$?curl?-s?-L?https://nvidia.github.io/nvidia-docker/gpgkey?|?sudo?apt-key?add?-
$?curl?-s?-L?https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list?|?sudo?tee?/etc/apt/sources.list.d/nvidia-docker.list

2. 安裝

$?sudo?apt-get?update
$?sudo?apt-get?install?-y?nvidia-docker2
$?sudo?systemctl?restart?docker

3. 驗(yàn)證

執(zhí)行以下命令：

$?docker?run?--rm?--gpus?all?nvidia/cuda:10.2-base?nvidia-smi?

如果輸出跟直接在宿主機(jī)上執(zhí)行?nvidia-smi?一致則說明安裝成功。如果跑的深度學(xué)習(xí)模型使用的是 tensorflow 可以在容器里執(zhí)行:

import?tensorflow?as?tf
tf.contrib.eager.num_gpus()

如果輸出了宿主機(jī)上的 Nvidia 顯卡數(shù)量，則模型能使用到顯卡加速。如果使用的是 pytorch 可以在容器里執(zhí)行：

import?torch
torch.cuda.is_available()

如果輸出 True 證明環(huán)境也成功了，可以使用顯卡。

4. 使用示例

• 使用所有顯卡

$?docker?run?--rm?--gpus?all?nvidia/cuda?nvidia-smi?
$?docker?run?--rm?--runtime=nvidia?-e?NVIDIA_VISIBLE_DEVICES=all?nvidia/cuda?nvidia-smi?

• 指明使用哪幾張卡

$?docker?run?--gpus?'"device=1,2"'?nvidia/cuda?nvidia-smi?
$?docker?run?--rm?--runtime=nvidia?-e?NVIDIA_VISIBLE_DEVICES=1,2?nvidia/cuda?nvidia-smi

到這里在 Docker 下使用 Nvidia 顯卡加速計(jì)算的基礎(chǔ)環(huán)境搭建就介紹完了。后續(xù)我們可以繼續(xù)研究一下 K8S 下調(diào)度 GPU 的實(shí)現(xiàn)。

本文轉(zhuǎn)載自：「lxkaka」，原文：https://lxkaka.wang/docker-nvidia/，版權(quán)歸原作者所有。