本篇文章将分享如何通过 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 。
如果你对如何封装基础镜像感兴趣,可以继续阅读本章节,如果你只关心如何快速跑模型,可以直接阅读下一章节内容。
言归正传,出于下面三个原因,我建议想在本地快速复现模型的同学采用容器方案 :
- 想要避免不同项目之间的环境造成干扰(污染)
- 想要确保项目依赖清晰,任何人都能够在任何设备上复现结果
- 想要复现模型的时间成本更低一些,不喜欢折腾 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 安装为例:
代码语言: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,在完成上面的基础环境安装之外,我们还需要指定 torch
和 torchvision
版本,关于这点在PyTorch 社区里曾有过 一些讨论。
RUN pip3 install --upgrade torch==1.9.0 torchvision==0.10.0
在上面的命令中,我们会将 torch
替换为指定版本。实际构建镜像的过程中,大概需要额外下载 800MB 的数据。即使我们使用了国内的软件源,时间可能也会比较漫长,可以考虑去冰箱里拿一罐冰可乐,缓解等待焦虑。