作为服务的Segment Anything编码器和解码器

快速入门

构建创建.marc文件的容器

将使用vit_h权重下载两个模型。

docker build -t sam-builder -f Dockerfile-build .

将.marc文件复制到主机以进行本地测试

docker run -d --name sam-builder1 sam-builder
docker cp sam-builder1:/home/model-store ./

我们将其复制到model-store，并通过GPU和CPU Torchserve容器在本地使用。

一旦复制了模型，就可以删除容器。

docker rm -f sam-builder1

构建用于图像编码的GPU torchserve容器

使用GPU，推理时间应为1.8秒或更少，具体取决于GPU。在较旧的1080 Ti Pascal GPU上，不进行编译的推理时间为1.67秒。

docker build -t sam-gpu -f Dockerfile-gpu .
bash start_serve_encode_gpu.sh

构建用于图像解码的CPU torchserve容器

docker build -t sam-cpu -f Dockerfile-cpu .
bash start_serve_decode_cpu.sh

在CPU上测试编码服务

默认情况下，CPU服务在7080上运行。GPU服务默认在8080上。

curl http://127.0.0.1:7080/predictions/sam_vit_h_encode -T ./data/sample-img-fox.jpg

测试

tests/中的所有测试都涵盖了解码器的功能。在test_decode.py中需要在不同测试之间共享输出时，在conftest.py中的pytest fixtures中运行decode.py中的逻辑。

要开始运行测试，请确保您有测试模型。您应该有与推理期间使用的相同模型，包括

(test.py3.10) (base) rave@rave-desktop:~/segment-anything-services/tests/models$ tree
.
├── sam_vit_h_4b8939.pth
└── sam_vit_h_decode.onnx

您可以通过解压缩从上一步中启动的sam-builder容器中复制到model-store的.marc文件来获取这两个模型。然后，将.onnx和.pth文件移动到./tests/models/

unzip model-store/sam_vit_h_decode.mar -d ./sam_decode_mar
unzip model-store/sam_vit_h_encode.mar -d ./sam_encode_mar
cp ./sam_decode_mar/sam_vit_h_decode.onnx ./tests/models/
cp ./sam_decode_mar/sam_vit_h_4b8939.pth ./tests/models/

使用安装测试环境（点击此处）： pip install hatch

然后，创建环境。我使用Python 3.10进行了测试，由于onnxruntime版本问题，Python 3.11无法工作。

hatch -e test.py3.10 shell

然后，使用pytest运行测试

pytest tests

无Docker的本地设置

1. 下载模型权重

如果您有权限，请从devseed s3下载

aws s3 sync s3://segment-anything/model-weights/ model-weights

否则，从原始仓库获取检查点：（点击此处）

2a. 打包torch权重以进行GPU编码

此步骤耗时较长，可能是因为未编译的权重很大。在后续步骤中，打包ONNX模型会更快。

mkdir -p model_store_encode
torch-model-archiver --model-name sam_vit_h_encode --version 1.0.0 --serialized-file model-weights/sam_vit_h_4b8939.pth --handler handler_encode.py
mv sam_vit_h_encode.mar model_store_encode/sam_vit_h_encode.mar

2b. 导出ONNX模型以进行CPU解码

mkdir -p models
python scripts/export_onnx_model.py --checkpoint model-weights/sam_vit_h_4b8939.pth --model-type vit_h --output models/sam_vit_h_decode.onnx

2c. 使用处理程序打包ONNX模型以进行CPU解码

我们将将其放在model_store_decode目录中，以保持ONNX模型文件与torchserve .mar模型存档区分开来。model_store/是由Torchserve在容器中自动创建的，这就是为什么我们在本地创建一个名为“model_store_decode”的文件夹。

mkdir -p model_store_decode
torch-model-archiver --model-name sam_vit_h_decode --version 1.0.0 --serialized-file models/sam_vit_h_decode.onnx --handler handler_decode.py
mv sam_vit_h_decode.mar model_store_decode/sam_vit_h_decode.mar

构建jupyter服务器容器

使用此容器在具有地理空间和pytorch依赖项的GPU启用jupyter笔记本服务器中测试模型。

docker build -t sam-dev -f Dockerfile-dev .

5. 测试端点

您可以通过运行test_endpoint.ipynb来使用您刚才启动的两个服务。此笔记本的依赖项很少，您可以选择自行安装，或者可以在下面的jupyter服务器中运行它们。

6. 运行jupyter服务器容器

这是一个具有GPU的容器，已设置SAM和一些我们常用的一些依赖项。您可以使用它在一个笔记本环境中尝试SAM模型。如果没有GPU，请删除--gpus参数。

docker run -it --rm \
    -v $HOME/.aws:/root/.aws \
    -v "$(pwd)":/segment-anything-services \
    -p 8888:8888 \
    -e AWS_PROFILE=devseed \
    --gpus all sam-dev

部署

1. 使用npm安装部署依赖项

npm install -g @openaddresses/deploy

2. 通过导航到AWS SSO登录页面并选择“命令行访问”来设置凭据，以复制临时凭据。将这些凭据粘贴到~/.aws/credentials中，并重命名AWS配置文件（向Ryan咨询）。
3. touch ~/.deployrc.json并填充以下内容

   {
    "profilename": {
        "region": "us-east-1"
    }
}

4. 运行deploy init以生成仓库的.deployrc.json配置。本地和全局配置都需要部署。根据配置文件名提示填写参数
5. 提交并提交包含更改的PR。等待所有GitHub操作完成，以便构建模型存档和Docker镜像
6. 运行deploy update prod将更改部署到生产环境

（可能）常见问题

问题：为什么有两个服务？

答案：我们正在探索在CPU解码器之外以按需方式运行图像编码的成本效益方法。最终我们希望删除后端CPU torserve的需求，并在浏览器中运行解码。

问题：我可以贡献或提问吗？

答案：这是一个更偏向于“公开工作”的仓库，我们希望与他人分享，而不是一个维护的项目。但如果您有任何想法，请随时提出问题。请理解如果我们不回应或响应缓慢。

许可证

模型和代码受Apache 2.0许可证许可。

参考文献

Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., ... Girshick, R. (2023). Segment Anything. arXiv:2304.02643. （点击此处）

scripts/export_onnx_model.ipynb和notebooks/sam_onnx_model_example_fox.ipynb来自原始仓库。