本文档提供了有关 Segment Anything Model (SAM) 仓库中提供的三种不同模型大小的技术信息、它们的架构差异、性能权衡和实现细节。有关自定义蒙版生成参数,请参阅自定义蒙版生成。
SAM 提供三种预训练模型变体,具有不同的尺寸-性能权衡,所有这些都基于 Vision Transformer (ViT) 架构
来源:segment_anything/build_sam.py46-52
模型变体之间的主要区别在于图像编码器架构参数,而提示编码器和蒙版解码器组件在所有变体中基本保持不变。
| 参数 | ViT-H (Huge) | ViT-L (Large) | ViT-B (Base) |
|---|---|---|---|
| 嵌入维度 | 1280 | 1024 | 768 |
| Transformer 深度 | 32 层 | 24 层 | 12 层 |
| 注意力头数 | 16 | 16 | 12 |
| 全局注意力索引 | [7, 15, 23, 31] | [5, 11, 17, 23] | [2, 5, 8, 11] |
| 参数数量 | 最大 | 中等 | 最小 |
| 检查点文件大小 | 约 2.4GB | 约 1.2GB | 约 375MB |
来源:segment_anything/build_sam.py14-44
此图显示了模型变体在代码库中的创建方式
来源:segment_anything/build_sam.py54-107
以下是加载不同模型变体的方法
加载的模型随后可以与 SamPredictor 或 SamAutomaticMaskGenerator 一起使用,就像使用任何其他 SAM 模型一样。
来源:README.md102-114 segment_anything/build_sam.py46-52
官方检查点文件在仓库中有所引用
选择模型变体时,请考虑这些性能特征
| 模型 | 相对质量 | 内存需求 | 推理速度 | 适用场景 |
|---|---|---|---|---|
| ViT-H | 最高 | 最高(约 2.5GB+ 模型,推理需要较高的 VRAM) | 最慢 | 高精度要求,提供强大的 GPU |
| ViT-L | 中高 | 中等(约 1.2GB 模型) | 中等 | 质量与性能的平衡 |
| ViT-B | 较低 | 最低(约 375MB 模型) | 最快 | 资源受限环境,实时应用 |
使用自动蒙版生成脚本时,使用 --model-type 参数指定模型类型
来源:scripts/amg.py41-46 scripts/amg.py196-201
在导出到 ONNX 进行部署时(尤其是在 Web/移动/边缘设备上),模型大小变得尤为重要
对于 Web 演示和资源受限的环境,推荐使用 ViT-B 模型,因为它具有更快的加载时间和更低的内存使用量,可以提供更好的用户体验。
模型大小的选择会影响
SAM 架构在所有三个模型变体中保持一致,只有图像编码器参数发生变化。这意味着使用任何一个模型的 API 都是相同的——唯一的区别在于初始化。