尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Docker部署SAM2图像分割模型:Linux环境与GPU透传实战指南

Docker部署SAM2图像分割模型:Linux环境与GPU透传实战指南 上周朋友让我帮他在Linux服务器上部署SAM2做图像分割我当时的第一反应是这东西装起来不麻烦吧。结果一上手CUDA版本、PyTorch、OpenCV、Hydra、tile这些依赖一个个冒出来装到一半系统里多了一堆乱七八糟的库最后折腾了一天才勉强跑起来。后来我换成Docker方案五分钟不到就部署好了模型推理也稳。这期就把这套“SAM2 Docker Linux”的完整部署流程拆开来讲从环境准备到labelme接入从GPU透传到常见坑位一条条过希望能帮你避开那些我踩过的坑。这个方案适合谁两类人最受益一是做图像分割、目标检测相关算法的同学想在Linux服务器上快速验证SAM2效果二是做标注工具集成的工程师想把SAM2能力嵌进labelme这类交互式标注工具里。无论你是刚接触Docker还是已经玩得很熟下面的操作都可以直接照抄关键参数和踩坑点我也会单独标出来。1. 为什么我坚持用Docker跑SAM21.1 原生部署SAM2到底麻烦在哪先还原一下裸机部署SAM2的典型场景。你打开官方仓库README上来就是PyTorch 2.3.1、torchvision、opencv-python、hydra-core、tile等一堆依赖然后你开始用pip装。装完以后发现PyTorch的CUDA版本和本机的CUDA driver对不上模型加载直接报CUDA error: no kernel image is available for execution on the device。就算你侥幸装好了后面还有更恶心的系统里的OpenCV版本和别的项目冲突改了一个库的版本又把另一个服务搞崩了。这种“环境地狱”在做深度学习项目时几乎人人都会遇到。SAM2本身依赖不算极端但它对PyTorch和CUDA的绑定比较紧而且不同版本的checkpoint对模型结构也有要求动不动就是一个库版本不匹配的问题。如果只是在开发机上自己玩忍一忍也就算了。可一旦要部署到客户的Linux服务器、要交付给运维、要换一台GPU机器重来一遍这种手动装环境的方式就是灾难。每台机器都当成“实验环境”去从零搭建本来就是不合理的。1.2 Docker怎么把“环境地狱”变成“拉箱子”Docker的核心思路可以理解成整屋搬家你把SAM2所需的操作系统库、Python依赖、CUDA runtime、模型代码全部冻结到一个镜像里哪个环境需要直接把整个“集装箱”拉过去运行。宿主机只需要装一个Docker引擎上面的库随便怎么乱都不影响别人容器内部运行什么版本都是自己说了算。这里有个关键点需要说清楚容器内的CUDA其实不是“显卡驱动”。Docker通过NVIDIA Container Toolkit把宿主机的GPU设备透传给容器容器内运行的是CUDA runtime和cuDNN而宿主机只需要提供GPU driver。所以哪怕容器里装的是CUDA 11.8宿主机驱动是535或更高版本只要能兼容就能正常跑。这种“runtime在容器内、driver在宿主机”的分离设计大大缓解了GPU环境的版本冲突问题。用Docker之后我部署SAM2的步骤就变成了写好或拉一个镜像跑容器挂载数据目录完事。换机器也一样把镜像导过去或重新build一次就行所有环境细节全部固化在镜像里。这就是我坚持用Docker的根本原因。2. 部署前的必备功课Docker安装与GPU透传2.1 安装Docker和几个高频命令既然目标是Linux环境我就不推荐Docker Desktop了。Docker Desktop更适合Windows、Mac上的本地开发在Linux服务器上直接用Docker Engine就行。安装方式我一般用官方脚本或系统包管理工具以Ubuntu为例sudo apt update sudo apt install docker.io -y sudo systemctl enable --now docker sudo docker version如果是CentOS / Rocky Linux用sudo yum install -y docker-ce或者docker包也行。装完以后把当前用户加进docker组省得每次都要sudosudo usermod -aG docker $USER newgrp docker验证一下docker ps能不能正常输出。这里经常会踩一个坑permission denied原因就是用户不在docker组里加进去之后重新登录终端就能解决。下面几个Docker命令是平时用得最多的docker ps -a # 查看所有容器 docker images # 查看本地镜像 docker exec -it 容器名 bash # 进入容器 docker logs 容器名 # 查看容器日志 docker rm -f 容器名 # 强制删除容器记不住没关系搜一下就能找到但高频的这几个最好心里有数后面跑容器的时候天天用。2.2 配置GPU支持NVIDIA Container ToolkitDocker本身不认识GPU需要装一个NVIDIA Container Toolkit把显卡设备映射进容器。安装方式可以参考NVIDIA官方文档Ubuntu上大致是这样distribution$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list sudo apt-get update sudo apt-get install -y nvidia-container-toolkit sudo nvidia-ctk runtime configure --runtimedocker sudo systemctl restart docker安装完成后先验证宿主机GPU驱动nvidia-smi再验证容器内能不能看到GPUdocker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi如果容器里能打印出类似宿主机的nvidia-smi信息说明GPU透传已经通了。这里有个特别常见的现象宿主机的nvidia-smi正常但容器里nvidia-smi: not found。其实不一定是你装错了有些精简镜像里根本没有nvidia-smi这个工具你只需要在容器里跑个Pythonimport torch; print(torch.cuda.is_available())确认即可。注意如果你的服务器没有NVIDIA显卡或者这一步只是想先跑通CPU流程可以跳过GPU部分。SAM2的CPU推理也能跑只是慢很多。对个人验证来说先用CPU跑一张图确认流程没问题再去上GPU是性价比最高的路线。3. 5分钟部署实操从拉镜像到跑通demo3.1 拿到SAM2代码和模型权重部署前需要准备两个东西官方代码仓库和预训练权重。SAM2的官方仓库在GitHub上直接clone下来git clone https://github.com/facebookresearch/segment-anything-2.git cd segment-anything-2模型权重通常可以从官方release页面或Hugging Face仓库下载有sam2_hiera_large.pt、sam2_hiera_base_plus.pt、sam2_hiera_small.pt等几个版本大小从几十MB到2.4GB不等。如果你只是测试先下small或base_plus就够了Large模型对显存要求更高。把权重文件放到一个管理清晰的目录比如~/models/后面启动容器时直接挂载进容器这样镜像本身不用重复烧录大文件。3.2 构建Docker镜像官方仓库里带了Dockerfile基于pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime。直接构建cd segment-anything-2 docker build -t sam2:dev .第一次构建会下载基础镜像和若干依赖耗时取决于网络。如果构建很慢也可以考虑自己写一个更精简的Dockerfile基础镜像换成pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime然后把仓库代码手动COPY进去安装依赖时只装推理所需的部分比如FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime WORKDIR /root RUN git clone https://github.com/facebookresearch/segment-anything-2.git \ cd segment-anything-2 \ pip install -e . -q \ pip install opencv-python-headless ENV PYTHONUNBUFFERED1 CMD [/bin/bash]这种方式能有效缩短构建时间。不过我个人的建议是第一次老老实实按官方Dockerfile构建构建过程本身就是一次很好的验证能提前暴露出依赖缺失的问题。3.3 启动容器并运行图像分割镜像构建好之后用docker run启动容器。这里的关键参数是GPU透传、目录挂载、共享内存大小docker run -it --rm \ --gpus all \ --shm-size8g \ -v ~/models:/root/models \ -v ~/data:/root/data \ -v ~/output:/root/output \ sam2:dev--shm-size8g是一个很容易被忽略的参数。PyTorch的DataLoader在多进程模式下会用到/dev/shm默认的64MB根本不够用跑着跑着就报Bus error或者Cannot allocate memory我之前就被这个坑折磨过。进入容器后先确认依赖和模型都正常cd /root/segment-anything-2 python -c import torch; print(torch.__version__, torch.cuda.is_available())下载好的权重文件放到~/models后容器内路径就是/root/models。现在可以写一个最简单的推理脚本用点提示point prompt来分割目标python - EOF import torch from PIL import Image import numpy as np import matplotlib.pyplot as plt from sam2.build_sam import build_sam2 from sam2.sam2_image_predictor import SAM2ImagePredictor checkpoint /root/models/sam2_hiera_base_plus.pt model_cfg configs/sam2/sam2_hiera_b.yaml sam2 build_sam2(model_cfg, checkpoint, devicecuda) predictor SAM2ImagePredictor(sam2) image Image.open(/root/data/cat.jpg).convert(RGB) image_np np.array(image) predictor.set_image(image_np) # 传入你要分割的物体上的一个点坐标和标签 input_point np.array([[500, 375]], dtypenp.float32) input_label np.array([1], dtypenp.int64) masks, scores, logits predictor.predict( point_coordsinput_point, point_labelsinput_label, multimask_outputTrue, ) # 取最高分mask保存 best_mask masks[np.argmax(scores)] plt.imsave(/root/output/mask.png, best_mask.astype(np.uint8) * 255) print(saved mask to /root/output/mask.png) EOF这段脚本核心就三步加载模型喂一张图传一个点坐标输出一个mask。我第一次在Docker容器里跑通时整个推理过程不到半分钟那种感觉确实爽。3.4 高频操作速查表操作命令构建SAM2镜像docker build -t sam2:dev .查看本地镜像docker images启动GPU容器docker run -it --rm --gpus all --shm-size8g -v ~/models:/root/models sam2:dev进入已运行容器docker exec -it 容器名 bash查看容器日志docker logs 容器名删除所有停止的容器docker container prune4. 进阶玩法把SAM2做成labelme的标注后端4.1 为什么需要把SAM2接入labelme很多人问“如何在labelme中部署sam2”其实真正的需求是标图的时候不想一个点一个点地抠轮廓想用SAM2的交互式分割能力辅助标注。你点一下物体内部SAM2就能生成完整mask标注效率翻倍。labelme本身是一个图像标注工具要接外部模型推理能力我的做法不是去改labelme的源码而是把SAM2封装成一个HTTP推理服务再用辅助脚本把mask结果转成labelme能读懂的JSON格式。这样解耦最干净模型挂了不影响主流程模型升级也只需要换容器。4.2 FastAPI包一层推理接口在SAM2容器里装FastAPI和uvicorn然后写一个最简推理服务# server.py import json import numpy as np from fastapi import FastAPI, UploadFile, File, Form from PIL import Image from sam2.build_sam import build_sam2 from sam2.sam2_image_predictor import SAM2ImagePredictor app FastAPI() sam2 build_sam2( configs/sam2/sam2_hiera_b.yaml, /root/models/sam2_hiera_base_plus.pt, devicecuda ) predictor SAM2ImagePredictor(sam2) app.post(/predict) async def predict( file: UploadFile File(...), points: str Form(...), # 例如 [[100,200],[300,400]] labels: str Form([[1],[1]]) ): image Image.open(file.file).convert(RGB) predictor.set_image(np.array(image)) pts np.array(json.loads(points), dtypenp.float32) labs np.array(json.loads(labels), dtypenp.int64) masks, scores, _ predictor.predict( point_coordspts, point_labelslabs, multimask_outputTrue ) best masks[np.argmax(scores)] mask_np best.astype(np.uint8) return {mask: mask_np.tolist(), score: float(np.max(scores))}启动服务uvicorn server:app --host 0.0.0.0 --port 8000注意一定要在容器启动时把端口映射出来-p 8000:8000。这个接口的设计很简单上传图片 点坐标返回mask的二维数组。4.3 在labelme里写脚本调用labelme里调用这个接口我的方案是一个独立的Python脚本先读取labelme导出的JSON或直接用labelme的API拿到当前图片路径和点击的点坐标调用上面的HTTP接口拿回mask后转成多边形写回labelme的JSON文件里。大致思路如下import json import requests import base64 def call_sam2(image_path, points): with open(image_path, rb) as f: files {file: f} data {points: json.dumps(points), labels: json.dumps([[1] for _ in points])} resp requests.post(http://localhost:8000/predict, filesfiles, datadata) return resp.json()[mask] def mask_to_polygon(mask): from skimage import measure contours measure.find_contours(mask, 0.5) # 取最大轮廓转成labelme的polygon格式 ... # 流程拿到图上点击的点 - call_sam2 - mask_to_polygon - 写入json[shapes]提示我这样接入能稳定跑通但它不是严格意义上的“labelme插件开发”。如果你想做成真正的插件可以在labelme的自定义标注插件框架里去实现一个后端适配器。对于一般项目使用独立脚本方案更简单、更好维护临时改个逻辑也不影响labelme本体。5. 常见问题与排查技巧实录5.1 Docker容器里看不到GPU这是出现频率最高的问题。排查步骤我建议按顺序来宿主机执行nvidia-smi看驱动是否正常。确认NVIDIA Container Toolkit已经安装sudo nvidia-ctk runtime configure --runtimedocker执行过。查看/etc/docker/daemon.json里面应该有类似下面的内容{ runtimes: { nvidia: { path: nvidia-container-runtime, runtimeArgs: [] } } }重启Dockersudo systemctl restart docker。大多数看不到GPU的情况问题都出在第2步或第3步尤其是改过daemon.json后忘了重启Docker怎么跑都不生效。5.2 OOM和显存不足怎么处理SAM2的显存占用和模型尺寸、输入图像尺寸直接相关。Large模型在1080p图片上推理显存占用可能到10GB以上。如果你的显卡只有8GB建议换sam2_hiera_base_plus.pt或sam2_hiera_small.pt同时控制输入图像长边不超过1024。Docker本身也可以限制容器内存docker run -it --rm -m 16g --gpus all sam2:dev另外注意PyTorch跑久了会有显存碎片化问题如果一个进程退出后显存没有释放可以用fuser -v /dev/nvidia*查看占用进程必要时kill掉。5.3 模型权重下载慢、时区问题和中文乱码权重文件动辄几百MB下载速度可能不稳定。我的建议是下载任务放到宿主机通过挂载目录让容器读取不要让容器每次启动都去拉取。这样就算换容器权重文件也不用重新下载。容器时区默认是UTC日志时间看起来会很别扭。启动时加参数-e TZAsia/Shanghai如果脚本里输出中文变成乱码进入容器后设置环境变量export LANGC.UTF-8或者在启动容器时加上-e LANGC.UTF-8。顺带说一句如果你的输入数据是在Windows环境下压缩再传到Linux的解压后文件名乱码那是zip编码问题不是Docker的问题需要用unzip -O GBK这种参数处理。5.4 常见问题速查表问题可能原因解决办法容器内无法使用GPUToolkit未配置运行sudo nvidia-ctk runtime configure --runtimedocker并重启启动容器报docker: Error response from daemon: could not select device driver缺少nvidia runtime重启Docker或安装NVIDIA Container Toolkit推理报CUDA out of memory显存不足或图片过大换小模型、限制图片尺寸、减小batch容器内中文显示乱码缺少UTF-8环境变量加-e LANGC.UTF-8容器时间不对默认UTC时区加-e TZAsia/Shanghai模型加载报KeyError: sam2_hiera_b模型配置和权重不匹配检查model_cfg和checkpoint是否对应6. Docker Compose管理复杂任务6.1 为什么需要compose当你的部署不只是“跑一个容器”而是需要同时启动SAM2推理服务、labelme辅助脚本、可能还有一个前端预览页再带上端口映射、环境变量、多个挂载目录用一长条docker run命令管理就会变得非常痛苦。这时候上Docker Compose最合适。Compose用YAML文件把容器的所有配置固化下来一条docker compose up -d就能把整套服务拉起来换机器部署也只需要复制配置文件。对团队协作来说这个文件就是所有人共享的环境说明书。6.2 一个可用的docker-compose.yml样例services: sam2: image: sam2:dev container_name: sam2_infer runtime: nvidia environment: - NVIDIA_VISIBLE_DEVICESall - TZAsia/Shanghai - LANGC.UTF-8 - PYTHONUNBUFFERED1 volumes: - ~/models:/root/models - ~/data:/root/data - ~/output:/root/output shm_size: 8g ports: - 8000:8000 command: uvicorn server:app --host 0.0.0.0 --port 8000这里一个比较新的改动是新版Docker Compose推荐直接用gpus: all而不是runtime: nvidiaservices: sam2: image: sam2:dev container_name: sam2_infer gpus: all shm_size: 8g volumes: - ~/models:/root/models - ~/data:/root/data ports: - 8000:8000两种方式选一种即可。启动docker compose up -d docker compose logs -f如果你需要把容器导出到没有外网的机器上还可以用docker save sam2:dev | gzip sam2.tar.gz打包到目标机器再gunzip -c sam2.tar.gz | docker load。不过模型权重文件因为体积大建议还是用挂载目录的方式镜像只存代码和依赖这样镜像体积会小很多。结尾前再讲两个实际体会第一个体会Docker部署SAM2最大的价值其实是“可重置”。以前我手动装环境一旦某个库跟系统库冲突就只能硬着头皮修修不好还得重装系统。现在Docker用完即焚容器删掉重建五分钟就回到初始状态试错成本非常低。这种“反脆弱”的特性对个人实验和团队协作都是极大的解放。第二个体会我在实际接入labelme的过程中发现模型服务的稳定性往往比推理速度更重要。我在FastAPI服务里加了单例加载模型、错误重试、超时控制之后才算是真正能在标注流程里稳定使用。所以别急着把模型、标注工具、界面全串起来先把推理服务本身调稳再一层层往上接这条路走起来会顺畅很多。最后再分享一个小技巧部署时养成把模型路径、数据目录、端口号、镜像版本全部写进环境变量或.env文件的习惯换机器只需要改一个文件不用到处找命令里的硬编码路径。我自己就是因为没这样做吃过几次“换机器改完代码但忘了改路径”的亏。
返回列表