快速分割一切

本项目基于 https://github.com/facebookresearch/segment-anything 的分支开发

相关博客文章：https://pytorch.org/blog/accelerating-generative-ai/

安装

步骤 1

获取最新的 PyTorch 每日构建版

例如：

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121

或

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu

安装说明因平台而异。请参阅网站 https://pytorch.org/

步骤 2

安装软件包

pip install git+https://github.com/pytorch-labs/segment-anything-fast.git

使用方法

该软件包可作为 segment-anything 的直接替代品。

例如，如果你现在使用 from segment_anything import sam_model_registry，你应该可以直接使用 from segment_anything_fast import sam_model_registry。

然而，你可能来这里是为了尝试快速推理版本。因此，我们还创建了 sam_model_fast_registry，它自动应用以下设置：

• 设置 eval 模式
• 使用 bfloat16
• 启用 torch.compile，并设置 max-autotune
• 使用自定义 Triton 内核，为长序列长度的相对位置编码实现 SDPA

特别是，自定义 Triton 内核是为 A100 编写的。如果你没有使用 A100，我们会尝试在你的设备上重新运行自动调优，并在本地保存最佳配置。 你可能仍会遇到性能问题，因此可以通过设置环境变量 SEGMENT_ANYTHING_FAST_USE_FLASH_4=0 来禁用该内核。

请注意，第一次运行此模型时，你可能需要等待一段时间以完成编译。

如果你想了解如何重现所有结果的详细信息，请查看上面实验文件夹中的 README。

如果你发现缺少功能或遇到问题，请不要犹豫，直接在 Github 上提出 issue。谢谢。

结果

结果展示了一系列技术的瀑布图。

从左到右，这些技术被组合使用。

这意味着最后一个柱状图是以下技术的组合：

• bfloat16
• torch.compile 与 max-autotune
• torch.scaled_dot_product_attention
• 自定义 Triton 内核，为长序列长度的相对位置编码实现 SDPA
• NestedTensors
• 动态 int8 对称量化
• 2:4 稀疏格式

许可证

segment-anything-fast 根据 Apache 2.0 许可证发布。