使用 Docker 来运行 HuggingFace 海量模型

2022-05-28 00:05:05 浏览数 (2)

本篇文章将分享如何通过 Docker 来在本地快速运行 Hugging Face 上的有趣模型。用比原项目更少的代码,和更短的时间成本将模型跑起来。

如果你熟悉 Python,那么绝大多数的模型项目,都可以用 10 分钟左右的时间,完成本地的部署和运行。

写在前面

为了方便展示,我选择了一个图像处理模型。在聊细节之前,让我们来一起看看这个模型项目的实际运行效果吧。

瞳孔和毛发细节也非常接近真实瞳孔和毛发细节也非常接近真实
相比较原始照片新照片清晰不少相比较原始照片新照片清晰不少

上面的图片处理所使用的 AI 模型,是我在 Hugging Face 上找到的。随着 Hugging Face 的爆火,平台上出现了越来越多的有趣的模型和数据集,目前仅模型数量就高达 4 万 5 千多个。

这些模型有一个有趣的特点,在云平台上跑的好好的,但是一旦想在本地跑起来就得各种“费劲”折腾,项目关联的 GitHub 中总是能看到用户反馈:这个模型和代码,我本地跑不起来,运行环境和调用代码搞起来太麻烦了。

Hugging Face 的数万开放模型Hugging Face 的数万开放模型

其实在日常的工作和学习中,我们也会经常遇到类似上面 Hugging Face 的情况:许多模型在“云上”跑的好好的,但是一到本地就跑不起来了,这或许是因为“操作系统环境、设备 CPU 架构(x86/ ARM)差异”、或许是因为“Python 运行时版本过高或过低”、或许是因为“某个 PIP 安装的软件包版本不对”、“冗长的示例代码中写死了一堆东西”...

那么,有没有什么偷懒的方法,可以让我们绕开这些浪费时间的问题呢?

在经过了一番折腾之后,我探索出了一个相对靠谱的方案:用 Docker 容器搭配 Towhee,制作模型的一键运行环境

譬如本文开头提到的这个模型,假如我们想进行快速调用,针对我们的图片进行一个快速修复处理,真的并不难:只需要一条 docker run 命令,搭配二、三十来行 Python 代码即可。

接下来,我就以腾讯 ARC 实验室开源的 GFPGAN 模型为例,来聊聊如何快速的把网上开放的模型快速的跑起来。

GitHub 上拥有两万颗星星的 GFPGANGitHub 上拥有两万颗星星的 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 ,以及编写过程中的思考:

考虑到模型可能需要在 x86ARM 两类设备上运行,推荐使用 miniconda3 这个基于 debian 内置了 conda 工具包的基础镜像。

代码语言:shell复制
FROM continuumio/miniconda3:4.11.0

关于基础环境镜像的使用,我推荐大家使用具体版本号,而不是 latest,这样可以让你的容器在需要重复构建的时候,也能保持“稳定”,减少“意外的惊喜”。如果你有特殊的版本需求,可以在这里找到更适合你的镜像版本。关于 condamini conda 相关的内容,本篇文章暂不赘述,感兴趣的同学可以从官方仓库中获得更多的信息。如果有需求的话,我会写一篇更详细的文章来聊聊它。

因为我们会频繁使用 OpenGL 的 API,所以我们需要在基础镜像中安装 libgl1-mesa-glx 软件包,如果你想了解这个软件包的详情,可以阅读 debian 官方软件仓库的文档,为了让安装时间更少,这里我调整了软件源为国内的“清华源”。

代码语言:shell复制
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 安装为例:

代码语言:shell复制
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 了。

考虑到大家的网络环境不相同,这里列出一些国内其他的常用镜像源。你可以根据你自己的情况,调整软件包下载源,来获取更快的软件包下载速度。

代码语言:shell复制
# 清华源
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,在完成上面的基础环境安装之外,我们还需要指定 torchtorchvision 版本,关于这点在PyTorch 社区里曾有过 一些讨论。

代码语言:shell复制
RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0

在上面的命令中,我们会将 torch 替换为指定版本。实际构建镜像的过程中,大概需要额外下载 800MB 的数据。即使我们使用了国内的软件源,时间可能也会比较漫长,可以考虑去冰箱里拿一罐冰可乐,缓解等待焦虑。

0 人点赞