Docker搭建本地AI模型(一)构建基础镜像,应用镜像,调用模型

容器化部署模型的优点:

  1. 避免不同项目之间的环境造成干扰(污染)
  2. 确保项目依赖清晰,任何人都能够在任何设备上复现结果
  3. 复现模型的时间成本更低,不喜欢折腾 80% 重复的模型调优之外的工作内容(尤其是环境、基础配置)

1.基础镜像

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

关于基础环境镜像的使用,推荐使用具体版本号,而不是 latest,这样可以让你的容器在需要重复构建的时候,也能保持“稳定”,减少意外的麻烦。

FROM continuumio/miniconda3:4.11.0

因为我们会频繁使用 OpenGL 的 API,所以我们需要在基础镜像中安装 libgl1-mesa-glx 软件包:

RUN apt install -y libgl1-mesa-glx

安装 pytorch

RUN conda install -y pytorch

为了让我们的基础镜像环境能够兼容 x86 和 ARM,在完成上面的基础环境安装之外,我们还需要指定 torchtorchvision 版本

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

为了后续运行各种 PyTorch 模型能够更省事,推荐直接在基础镜像中安装 Towhee:

RUN pip install towhee

一个基于 PyTorch 的模型使用的通用 Docker 基础镜像的 Dockerfile 就编写完毕了,完整文件内容

FROM continuumio/miniconda3:4.11.0

# 国内换源的代码加在此处
RUN apt install -y libgl1-mesa-glx

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 zhumao/docker-pytorch-base .

等到命令执行完毕,我们的 PyTorch 基础镜像就构建完成了。


2.使用 Python 编写模型调用程序

使用 Towhee ,并且额外实现一个小功能:扫描工作目录的所有图片,然后将他们分别交给模型去处理,最后生成一个静态页面,来将处理前后的图片进行对比展示。

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'))
)

3.具体模型使用的应用镜像

有了上文中的基础镜像之后,就只需要针对每个不同的模型做一些镜像依赖微调即可。

可以节约大量的镜像构建时间,以及本地磁盘空间。

接下来针对 GFPGAN 项目做应用镜像定制:

FROM zhumao/docker-pytorch-base:2023.11.19

我们需要在要制作的应用镜像中放置我们要使用的模型文件,以及完成相关 Python 依赖的补充下载。

至于具体的模型使用方式,不论是打包到镜像里,或者选择在使用的过程中动态的挂载,都是可以的。

在 GFPGAN 项目中,一共依赖两个模型文件,一个是 https://github.com/xinntao/facexlib 项目中基于 ResNet50 的人脸检测模型,另一个是用于图片修复的 GFPGAN 对抗网络模型。

第一个模型文件 detection_Resnet50_Final.pth ,我们可以在这里获取https://github.com/xinntao/facexlib/releases/tag/v0.1.0

第二个模型则需要我们根据自己的设备状况,来做具体选择:

•如果你需要使用 CPU 来跑模型,可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.2.0 中下载 GFPGANCleanv1-NoCE-C2.pth;或者在https://github.com/TencentARC/GFPGAN/releases/tag/v1.3.0 中下载 GFPGANv1.3.pth,这类模型可以完成黑白人像图片的处理。

•如果你可以使用 GPU 来跑模型,可以在 https://github.com/TencentARC/GFPGAN/releases/tag/v0.1.0 中下载 GFPGANv1.pth;或者在 https://share.weiyun.com/ShYoCCoc 中进行模型文件下载,这类模型可以处理带颜色的人像图片。

•除了 GitHub 之外,我们也可以选择直接从 Hugging Face 下载模型(只是可选版本不像上面那么多):

https://huggingface.co/TencentARC/GFPGANv1/tree/main

将下载好的模型文件放置于相同目录之后,继续完善 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 zhumao/docker-pytorch-base:2023.11.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 .


4.模型应用镜像的使用

如果上一步你已经下载了模型文件,并将模型文件打包到了镜像中,那么我们只需要下载一些黑白或者彩色的包含人像的图片(根据模型来选择),将它们放在一个目录中(比如 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 容器基础镜像、如何封装具体模型的应用镜像、如何快速的调用模型


国内换源加速:

写在Dockerfile里:

apt 和 pip 换源

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 pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

更多国内常用镜像源:

# 清华源
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

如何接入好看的Gradio图形界面:


参考:

https://cloud.tencent.com/developer/article/2091469

pytorch社区讨论: https://discuss.pytorch.org/t/failed-to-load-image-python-extension-could-not-find-module/140278/14

分享你的喜爱

留下评论

您的电子邮箱地址不会被公开。 必填项已用*标注