win11环境下本地部署Segment Anything示例
前言
Segment Anything作为一种效果优秀的算法受到广泛关注。笔者今天用了两三个小时的时间根据官方教程在本地进行了部署尝试,具体而言实现了使用官方提供的模型(vit_h)分割任意图片的效果。整体比较顺利,现将过程记录如下。
运行环境
笔者采用windows11系统,在英伟达GPU上使用官方自带jupyter代码运行。
安装
python环境
为不影响其他程序正常运行,笔者使用python环境管理工具anaconda新建名为“sam”的独立python环境专门用于运行SAM相关的代码。具体而言,安装完anaconda后,先在cmd命令界面运行“conda”指令确认是否将其添加入环境变量,确认已添加后开始新建名为“sam”的python环境。根据官网要求,python版本需大于等于3.8,故使用python3.8:
conda create -n sam python=3.8
新建完成后激活:
conda activate sam
pytorch环境
接下来安装pytorch,确保本地已安装cuda和cudnn等cuda基础环境,通过NVIDIA控制面板菜单栏的“帮助”->“系统信息”->“组件”的第三行“NVCUDA.DLL”查看CUDA版本。
笔者CUDA版本为12.2,在pytorch官网无对应版本,经测试按CUDA11.8来安装也能正常运行。官网教程要求pytorch大于等于1.7且torchvision大于等于0.8,故采用在激活的sam环境中使用pip命令安装最新版pytorch:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
安装完成后测试cuda是否可以正常运行,返回True说明pytorch安装成功。
安装Segment Anything
可采用git方式直接下载,但windows对git命令并不友好,故采用手动下载到本地的方式,网址为https://github.com/facebookresearch/segment-anything。
下载完成后将压缩包解压得到程序代码。为方便操作,可将cmd命令行的路径通过cd命令改为程序根目录。
根据官方教程安装Segment Anything。
接下来安装运行依赖包。
pip install opencv-python matplotlib onnxruntime onnx
官方教程还安装了pycocotools用于处理COCO格式数据,实际安装中提示缺乏其他依赖,考虑到本次仅测试此算法图像分割效果,故暂不安装此依赖。
放置Model Checkpoints
官方提供了训练好的Model Checkpoints用于测试,故可免去训练步骤直接测试效果。官网提供了三个版本,默认的ViT-H模型,以及ViT-L和ViT-B模型。笔者使用效果较好的ViT-H模型用于测试。
通过链接下载到本地,为不改动官方代码,临时将此模型放到程序包中的“notebooks”目录下,使程序可直接访问。也可放到其他位置,并将路径添加到程序中。
配置jupyter环境
由于本次使用jupyter notebook运行测试程序,故需额外安装jupyter。
pip install jupyter
在sam环境中启动jupyter。
jupyter notebook
运行此命令后浏览器会自动弹出jupyter notebook页面,路径默认为根目录,进入所下载的Segment Anything程序所在目录,打开“notebooks”文件夹中名为“automatic_mask_generator_example”的文件,选择刚才安装好依赖环境的“sam”内核,若只存在Python3内核,则可通过下述方式解决。
若只有默认的Python3内核,则切回cmd命令行页面,在激活的sam环境下安装内核。
首先使用conda安装ipykernel。
conda install ipykernel
在虚拟环境sam中创建kernel文件。
conda install -n sam ipykernel
将sam环境写入notebook的kernel中。
python -m ipykernel install --user --name sam --display-name sam-kernel
以上命令运行完成后回到jupyter notebook页面,刷新页面即可找到新建的sam-kernel内核。
至此,准备工作全部完成。
测试
刚才打开的“automatic_mask_generator_example”文件即为用于测试的程序环境,确保所用内核为sam环境。
通过快捷键Shift+Enter从上向下运行所有程序。
通过Example image小节可知默认将dog.jpg用于测试。
经过自动mask生成,可得到分割后的效果,识别效果还是很不错的。
此外,提供了另一分割方式以在较小的对象上获得更好的分割效果。
上下对比可以看出分割更加细致。
资源使用方面,GPU显存占用9G左右,远小于GPT等大模型,且官方同时提供资源占用较小的其他可用模型,便于本地部署和推广。
笔者从网上找了几张图片进行测试。
Segment Anything一大亮点是引入了prompt,可根据需要自定义分割任务。通用大模型是AI发展的趋势,期待更多领域、泛化能力更强的通用大模型被提出。