使用 Docker 來執行 HuggingFace 海量模型

語言: CN / TW / HK
時間 2022-05-20 16:42:00

本篇文章將分享如何通過 Docker 來在本地快速執行 Hugging Face 上的有趣模型。用比原專案更少的程式碼，和更短的時間成本將模型跑起來。

如果你熟悉 Python，那麼絕大多數的模型專案，都可以用 10 分鐘左右的時間，完成本地的部署和執行。

寫在前面

為了方便展示，我選擇了一個影象處理模型。在聊細節之前，讓我們來一起看看這個模型專案的實際執行效果吧。

上面的圖片處理所使用的 AI 模型，是我在 Hugging Face 上找到的。隨著 Hugging Face 的爆火，平臺上出現了越來越多的有趣的模型和資料集，目前僅模型數量就高達 4 萬 5 千多個。

這些模型 有一個有趣的特點 ，在雲平臺上跑的好好的，但是一旦想在本地跑起來就得各種“費勁”折騰，專案關聯的 GitHub 中總是能看到使用者反饋：這個模型和程式碼，我本地跑不起來，執行環境和呼叫程式碼搞起來太麻煩了。

其實在日常的工作和學習中，我們也會經常遇到類似上面 Hugging Face 的情況：許多模型在“雲上”跑的好好的，但是一到本地就跑不起來了，這或許是因為“作業系統環境、裝置 CPU 架構（x86/ ARM）差異”、或許是因為“Python 執行時版本過高或過低”、或許是因為“某個 PIP 安裝的軟體包版本不對”、“冗長的示例程式碼中寫死了一堆東西”…

那麼，有沒有什麼偷懶的方法，可以讓我們繞開這些浪費時間的問題呢？

在經過了一番折騰之後，我探索出了一個相對靠譜的方案：用 Docker 容器搭配 Towhee，製作模型的 一鍵執行環境 。

譬如本文開頭提到的這個模型，假如我們想進行快速呼叫，針對我們的圖片進行一個快速修復處理，真的並不難：只需要一條 docker run 命令，搭配二、三十來行 Python 程式碼即可。

接下來，我就以騰訊 ARC 實驗室開源的 GFPGAN 模型為例，來聊聊如何快速的把網上開放的模型快速的跑起來。

因為該模型基於 PyTorch，所以本篇文章裡，我們先聊聊如何製作基於 PyTorch 的模型使用的通用 Docker 基礎映象。如果同學們有需求，我就再聊聊其他模型框架。

製作 PyTorch 模型使用的通用 Docker 基礎映象

本章節的完整示例程式碼，我已經上傳到了 GitHub： https://github.com/soulteary/docker-pytorch-playground ，感興趣的同學可以自取。如果你想進一步省事，還可以直接使用我已經構建好的映象，來作為基礎映象使用： https://hub.docker.com/r/soulteary/docker-pytorch-playground 。

如果你對如何封裝基礎映象感興趣，可以繼續閱讀本章節，如果你只關心如何快速跑模型，可以直接閱讀下一章節內容。

言歸正傳，出於下面三個原因， 我建議想在本地快速復現模型的同學採用容器方案 ：

  1. 想要避免不同專案之間的環境造成干擾（汙染）
  2. 想要確保專案依賴清晰，任何人都能夠在任何裝置上覆現結果
  3. 想要復現模型的時間成本更低一些，不喜歡折騰 80% 重複的模型調優之外的工作內容（尤其是環境、基礎配置）

在瞭解到容器方案的優勢之後。接下來，我們來聊聊如何編寫這類基礎映象的 Dockerfile ，以及編寫過程中的思考：

考慮到模型可能需要在 x86 和 ARM 兩類裝置上執行，推薦使用 miniconda3 這個基於 debian 內建了 conda 工具包的基礎映象。 

FROM continuumio/miniconda3:4.11.0

關於基礎環境映象的使用，我推薦大家使用具體版本號，而不是 latest ，這樣可以讓你的容器在需要重複構建的時候，也能保持“穩定”，減少“意外的驚喜”。如果你有特殊的版本需求，可以在 這裡 找到更適合你的映象版本。關於 conda 和 mini conda 相關的內容，本篇文章暫不贅述，感興趣的同學可以從 官方倉庫 中獲得更多的資訊。如果有需求的話，我會寫一篇更詳細的文章來聊聊它。

因為我們會頻繁使用 OpenGL 的 API，所以我們需要在基礎映象中安裝 libgl1-mesa-glx 軟體包，如果你想了解這個軟體包的詳情，可以閱讀 debian 官方軟體倉庫的文件 ，為了讓安裝時間更少，這裡我調整了軟體源為國內的“清華源”。 

RUN sed -i -e "s/deb.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && \
    sed -i -e "s/security.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && \
    apt update
RUN apt install -y libgl1-mesa-glx

當我們完成了基礎系統依賴庫的安裝之後，就可以開始準備模型執行環境了，以 PyTorch 安裝為例：

RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
RUN conda install -y pytorch

同樣為了節約 Python PyPi 軟體包的下載時間，我同樣將下載源切換到了國內的“清華源”。當 conda install -y pytorch 命令執行完畢之後，我們的基礎的執行環境也就 OK 了。

考慮到大家的網路環境不相同，這裡列出一些國內其他的常用映象源。你可以根據你自己的情況，調整軟體包下載源，來獲取更快的軟體包下載速度。

# 清華源
https://pypi.tuna.tsinghua.edu.cn/simple
# 阿里雲
http://mirrors.aliyun.com/pypi/simple
# 百度
https://mirror.baidu.com/pypi/simple
# 中科大
https://pypi.mirrors.ustc.edu.cn/simple
# 豆瓣
http://pypi.douban.com/simple

在上面的步驟中，我們大概需要下載接近 200MB 的軟體包（ conda 14MB、 pytorch 44MB、 mkl 140MB），需要有一些耐心。

為了讓我們的基礎映象環境能夠相容 x86 和 ARM，在完成上面的基礎環境安裝之外，我們還需要指定 torch 和 torchvision 版本，關於這點在PyTorch 社群裡曾有過 一些討論 。 

RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0

在上面的命令中，我們會將 torch 替換為指定版本。實際構建映象的過程中，大概需要額外下載 800MB 的資料。即使我們使用了國內的軟體源，時間可能也會比較漫長，可以考慮去冰箱裡拿一罐冰可樂，緩解等待焦慮。

在處理完上面的各種依賴之後，我們就來到了構建映象的最後一步。為了後續執行各種 PyTorch 模型能夠更省事，推薦直接在基礎映象中安裝 Towhee：

# https://docs.towhee.io/Getting%20Started/quick-start/
RUN pip install towhee

至此，一個基於 PyTorch 的模型使用的通用 Docker 基礎映象的 Dockerfile 就編寫完畢啦，為了方便閱讀，我在這裡貼出完整檔案內容：

FROM continuumio/miniconda3:4.11.0

RUN sed -i -e "s/deb.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && \
    sed -i -e "s/security.debian.org/mirrors.tuna.tsinghua.edu.cn/" /etc/apt/sources.list && \
    apt update
RUN apt install -y libgl1-mesa-glx

RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
RUN conda install -y pytorch

RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0

RUN pip install towhee

將上面的內容儲存為 Dockerfile 後，執行 docker build -t soulteary/docker-pytorch-playground . ，等到命令執行完畢，我們的 PyTorch 基礎映象就構建完成了。

如果你不想浪費時間構建，也可以直接使用我已經構建好的基礎映象（支援自動區分 x86 / ARM 架構裝置），直接從 DockerHub 下載即可：

# 可以直接下載最新版本
docker pull soulteary/docker-pytorch-playground
# 也可以使用帶有具體版本的映象
docker pull soulteary/docker-pytorch-playground:2022.05.19

搞定基礎映象之後，我們就可以繼續折騰前文提到的具體模型的執行環境和程式啦。

使用 Python 編寫模型呼叫程式

我們在 GFPGAN 專案中可以找到官方的模型使用示例： https://github.com/TencentARC/GFPGAN/blob/master/inference_gfpgan.py ，原始檔案比較長，大概有 155 行，這裡我就不貼了。

我在上一小節中提到過，我們可以使用 Towhee 來“偷懶”，比如可以將示例程式碼的行數縮短到 30 來行，並且額外實現一個小功能：掃描工作目錄的所有圖片，然後將他們分別交給模型去處理，最後生成一個靜態頁面，來將處理前後的圖片進行對比展示。

import warnings
warnings.warn('The unoptimized RealESRGAN is very slow on CPU. We do not use it. '
              'If you really want to use it, please modify the corresponding codes.')

from gfpgan import GFPGANer
import towhee

@towhee.register
class GFPGANerOp:

    def __init__(self,
                 model_path='/GFPGAN.pth',
                 upscale=2,
                 arch='clean',
                 channel_multiplier=2,
                 bg_upsampler=None) -> None:
        self._restorer = GFPGANer(model_path, upscale, arch, channel_multiplier, bg_upsampler)

    def __call__(self, img):
        cropped_faces, restored_faces, restored_img = self._restorer.enhance(
            img, has_aligned=False, only_center_face=False, paste_back=True)

        return restored_faces[0][:, :, ::-1]

(
    towhee.glob['path']('*.jpg') 
        .image_load['path', 'img']() 
        .GFPGANerOp['img','face']() 
        .show(formatter=dict(img='image', face='image'))
)

如果將上面為了保持“原汁原味”的 warnings 剔除掉，其實還可以獲得更短的行數。將上面的內容儲存為 app.py ，我們稍後使用。

搞定了呼叫模型所需要的程式之後，我們繼續來聊聊，如何製作具體模型（GFPGAN）執行所需要的應用容器映象。

製作具體模型使用的應用映象

這部分的完整程式碼，我同樣上傳到了 GitHub，方便大家“偷懶”： https://github.com/soulteary/docker-gfpgan 。配套的預構建映象在這裡 https://hub.docker.com/r/soulteary/docker-gfpgan 。

言歸正傳， 有了上文中的基礎映象之後，我們在日常玩耍的過程中，就只需要針對每個不同的模型做一些映象依賴微調即可 。

下面就來看看如何針對上文中提到的 GFPGAN 專案做應用映象定製吧。

同樣是以編寫 Dockerfile 為例，先來宣告我們正在構建的應用映象是基於上面的基礎映象。 

FROM soulteary/docker-pytorch-playground:2022.05.19

這樣做的好處是，在後續的日常使用中，我們可以節約大量的映象構建時間，以及本地磁碟空間。不得不說，模型類大容器特別能夠享受 Docker 特性帶來的便利。

接下來，我們需要在要製作的應用映象中放置我們要使用的模型檔案，以及完成相關 Python 依賴的補充下載。

考慮到國內網路下載 Hugging Face 和 GitHub 模型比較慢，還容易出現網路中斷。我推薦大家在做應用模型構建的時候，可以考慮提前進行依賴模型的下載，在構建映象的過程中，將模型放置到合適的目錄位置即可。至於具體的模型使用方式，不論是打包到映象裡，或者選擇在使用的過程中動態的掛載，其實都是可以的。

在 GFPGAN 專案中，我們一共依賴倆模型檔案，一個是 https://github.com/xinntao/facexlib 專案中基於 ResNet50 的人臉檢測模型，另一個是用於圖片修復的 GFPGAN 對抗網路模型，也就是傳統意義上的“主角”。

第一個模型檔案 detection_Resnet50_Final.pth ，我們可以在 https://github.com/xinntao/facexlib/releases/tag/v0.1.0 中獲取；第二個模型則需要我們根據自己的裝置狀況，來做具體選擇：

將下載好的模型檔案和新的 Dockerfile 檔案放置於相同目錄之後，我們來繼續完善 Dockerfile 的內容，完成專案依賴的安裝，並將模型放置在容器內合適的目錄位置： 

# 安裝模型相關程式碼庫
RUN pip install gfpgan realesrgan
# 將提前下載好的模型複製到指定位置，避免構建映象過程中的意外
COPY detection_Resnet50_Final.pth /opt/conda/lib/python3.9/site-packages/facexlib/weights/detection_Resnet50_Final.pth

# 根據你下載的模型版本做選擇，選一個模型檔案就行
COPY GFPGANCleanv1-NoCE-C2.pth /GFPGAN.pth
# COPY GFPGANCleanv1-NoCE-C2_original.pth /GFPGAN.pth
# COPY GFPGANv1.pth /GFPGAN.pth
# COPY GFPGANv1.3.pth /GFPGAN.pth

上面除了 gfpgan 之外，我還安裝了 realesrgan ，這個軟體包可以讓處理完畢的圖片中的人臉之外的背景也顯得更好看、更自然一些。

完成了基礎依賴、模型的配置之後，最後就是一些簡單的收尾工作了：

# 將上一步儲存的呼叫模型的程式拷貝到映象中
COPY app.py /entrypoint.py

# 宣告一個乾淨的工作目錄
WORKDIR /data
# 這裡可以考慮直接將我們要測試的資料集扔到容器裡
# 也可以考慮在執行過程中動態的掛載進去
# COPY imgs/*.jpg ./

# 補充安裝一些專案需要的其他依賴
RUN pip install IPython pandas
# 因為 Towhee 目前只支援直接展示模型結果
# 暫時還不支援將展示結果儲存為檔案
# 所以這裡需要打個小補丁，讓它支援這個功能
RUN sed -i -e "s/display(HTML(table))/with open('result.html', 'w') as file:\n            file.write(HTML(table).data)/" /opt/conda/lib/python3.9/site-packages/towhee/functional/mixins/display.py
CMD ["python3", "/entrypoint.py"]

上面的程式碼中，我添加了不少註釋來解釋每一步要做什麼，就不多贅述啦。額外解釋一下這裡的設計和思考，把上文中的 app.py 挪到 / 根目錄，而不是扔到工作目錄可以讓我們的程式在使用過程中更簡單，因為我計劃將工作目錄作為圖片的讀取和處理結果的儲存目錄。容器最後使用 CMD 而不是 ENTRYPOINT 來執行預設命令，也更方便使用者直接呼叫命令，或者進入容器除錯。

同樣的，為了方便閱讀，我將上面的 Dockerfile 內容合併到一起：

FROM soulteary/docker-pytorch-playground:2022.05.19

RUN pip install gfpgan realesrgan
COPY detection_Resnet50_Final.pth /opt/conda/lib/python3.9/site-packages/facexlib/weights/detection_Resnet50_Final.pth

# 尺寸大一些的模型檔案，可以選擇使用掛載的方式
# 而不在此處直接 COPY 到容器內部
COPY GFPGANCleanv1-NoCE-C2.pth /GFPGAN.pth

COPY app.py /entrypoint.py
WORKDIR /data
RUN pip install IPython pandas
RUN sed -i -e "s/display(HTML(table))/with open('result.html', 'w') as file:\n            file.write(HTML(table).data)/" /opt/conda/lib/python3.9/site-packages/towhee/functional/mixins/display.py
CMD ["python3", "/entrypoint.py"]

將上面的內容儲存為 Dockerfile 之後，我們執行命令，來完成應用映象的構建： 

docker build -t pytorch-playground-gfpgan -f Dockerfile .

片刻之後，我們就得到一個包含了模型和模型執行程式的應用映象啦。

接下來，我們來看看如何使用這個映象，來得到文章一開始時的模型執行結果。

模型應用映象的使用

如果上一步你已經下載了模型檔案，並將模型檔案打包到了映象中，那麼我們只需要下載一些黑白或者彩色的包含人像的圖片（根據模型來選擇），將它們放在一個目錄中（比如 data 目錄），然後執行一行命令就能夠完成模型的呼叫啦： 

docker run --rm -it -v `pwd`/data:/data soulteary/docker-gfpgan

如果你不願意費事找圖片，也可以直接使用我在專案中準備的示例圖片： https://github.com/soulteary/docker-gfpgan/tree/main/data 。

上面是針對應用映象中包含模型的情況，下面我們來看看如果應用映象中不包含模型要怎麼處理。

如果在上文構建應用模型映象時，沒有選擇將 GFPGAN 模型打包到映象中，那麼我們就需要使用檔案掛載的方式，來執行模型了。為了專案結構的清晰，我在專案中建立了一個名為 model 的目錄，來存放上文中提到的模型檔案。

完整的目錄結構類似下面這樣：

.
├── data
│   ├── Audrey\ Hepburn.jpg
│   ├── Bruce\ Lee.jpg
│   ├── Edison.jpg
│   ├── Einstein.jpg
│   └── Lu\ Xun.jpg
└── model
    └── GFPGANCleanv1-NoCE-C2.pth

當準備好模型和要處理的圖片之後，我們還是執行一條簡單的命令，來將檔案掛載到容器中，讓模型發揮“魔力”：

docker run --rm -it -v `pwd`/model/GFPGANCleanv1-NoCE-C2.pth:/GFPGAN.pth -v `pwd`/data:/data soulteary/docker-gfpgan

當命令執行完畢之後，在 data 目錄中，會多出一個 result.html 檔案，裡面記錄了模型處理前後的圖片結果。使用瀏覽器直接開啟，可以看到類似下面的結果：

寫到這裡，如何封裝 PyTorch 容器基礎映象、如何封裝具體模型的應用映象、如何快速的呼叫模型就都介紹完啦。如果後面有機會，我會聊聊如何基於這些映象做進一步的效能調優，以及聊聊 PyTorch 之外的映象封裝。

最後

本篇內容的完成，需要感謝兩位好朋友、Towhee 專案的核心開發者 @侯傑 、@郭人通 的幫助。解決了對於我這個 Python 菜鳥來說，雖然行數不多，但是非常麻煩的模型呼叫的內容。

下一篇相關的內容中，我計劃聊聊如何在 M1 裝置上進行模型訓練和推理，以及繼續實踐一些更有趣的 AI 專案。

–EOF

「」
「」