segmentation_models.pytorch
PyTorch图像语义分割框架,支持12种分割架构与500+预训练编码器,零门槛三行代码构建分割模型
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
PyTorch图像语义分割框架,支持12种分割架构与500+预训练编码器,零门槛三行代码构建分割模型
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。

图1:Segmentation Models PyTorch 项目标志
你可能见过这样的场景:医生对着肺部CT片用红笔一圈一划,标出可能的病灶区域;工程师在航拍照片上用不同颜色标记出道路、建筑、绿地;自动驾驶汽车实时"看到"路面的车道线、行人、障碍物。这些工作,本质上都是在做一件事——图像语义分割(Semantic Segmentation)。
过去,要从零构建一套图像分割模型,需要深厚的深度学习功底:设计网络架构、编写解码器逻辑、处理预训练权重迁移、适配各种 backbone……每一步都可能让人望而却步。Pavel Iakubovskii用一个Python库把这些工作全部打包好了——这就是 segmentation_models_pytorch(简称SMP),GitHub星标11,580,PyTorch生态中最受欢迎的图像分割工具库之一。
SMP的诞生并非偶然。作为一名计算机视觉工程师,Iakubovskii在实际项目中频繁遇到一个问题:现有的分割模型实现各自为政,Unet、DeepLabV3、FPN……每个架构都有独立的代码库,彼此之间几乎没有可复用性。如果想对比不同架构在同一数据集上的表现,需要翻阅大量不同的源码,适配不同的预处理逻辑。
2019年左右,他决定把这件事系统化——做一个统一的框架,让任何人都能以相同的方式调用不同的分割架构,只需要切换一个参数。项目在GitHub上公开后,迅速吸引了大量关注,目前已有超过1,800个fork和来自全球开发者的持续贡献。
SMP的设计哲学可以用一句话概括:一个解码器,多种编码器。
图像分割任务通常由两个核心部分组成:
Encoder(编码器):负责从输入图像中提取特征,类似人类的视觉皮层处理原始像素。在SMP中,Encoder实际上就是ImageNet预训练的分类网络(如ResNet、EfficientNet、ConvNeXt),只不过提取的是中间层的特征图(feature maps),而非最终分类结果。
Decoder(解码器):负责将Encoder提取的特征"翻译"为像素级的分割掩码(segmentation mask),即对图像中每个像素标注其所属类别。
这种设计带来了极大的灵活性。SMP支持12种主流分割架构(Unet、Unet++、DeepLabV3、DeepLabV3+、FPN、PSPNet、PAN、Linknet、MAnet、UPerNet、Segformer、DPT)和500+种预训练编码器(来自torchvision、timm等主流库)。你可以像组装乐高一样,自由组合不同架构和编码器,快速验证哪种组合最适合你的任务。

图2:SMP架构文档示意
对于没有分割模型经验的开发者,SMP的上手门槛极低。以下是一个完整的图像分割模型创建代码:
import segmentation_models_pytorch as smp
# 第一步:选择架构和预训练编码器
model = smp.Unet(
encoder_name="resnet34", # 编码器:ImageNet预训练的ResNet34
encoder_weights="imagenet", # 使用ImageNet预训练权重
in_channels=3, # RGB三通道
classes=3, # 输出3类分割结果
)
# 第二步:加载数据并训练
# (你的数据和训练循环代码)
# 第三步:推理预测
import torch
x = torch.randn(1, 3, 256, 256)
pred = model(x)
print(pred.shape) # torch.Size([1, 3, 256, 256])
只需三行核心代码,你就能得到一个基于ImageNet预训练ResNet34作为编码器的U-Net分割模型。如果你有GPU,整个模型创建和推理几乎是即时的。
项目还提供了大量Google Colab教程Notebook,涵盖二分类分割、多分类分割、预训练模型推理、模型导出ONNX等常见场景。不需要任何本地配置,直接在浏览器中运行即可体验。
SMP的另一大亮点是预训练权重的管理。项目在HuggingFace Hub上维护了专门的模型集合(如UPerNet、Segformer、DPT系列),支持直接通过timm库加载最新的预训练编码器权重。这意味着你可以直接使用EfficientNet-B7、ConvNeXt-Base、Swin Transformer等最新架构,而无需手动下载权重文件。
预训练权重还内置了与ImageNet训练时一致的预处理逻辑(归一化参数、通道顺序等),通过get_preprocessing_fn()即可获取对应的预处理函数,确保模型在推理时能以正确的方式处理输入图像。
SMP的代码结构非常清晰,核心模块组织如下:
segmentation_models_pytorch/base/:基础模型类和权重加载逻辑segmentation_models_pytorch/encoders/:原生支持的编码器实现(ResNet、DenseNet等),以及对timm库的集成segmentation_models_pytorch/decoders/:12种分割解码器的完整实现,每个子目录对应一种架构segmentation_models_pytorch/losses/:丰富的损失函数集合(Dice Loss、Focal Loss、Boundary Loss等),覆盖常用分割评估指标segmentation_models_pytorch/metrics/:IoU、Dice、F1等分割专用指标计算segmentation_models_pytorch/datasets/:常见数据集的封装(Carvana、Potsdam等)项目采用pyproject.toml管理依赖,使用pytest + pytest-cov + pytest-xdist进行测试,ruff做代码规范检查,文档由Sphinx生成,托管在ReadTheDocs。整体代码质量评分较高,测试覆盖率完整。
SMP主要面向以下场景:
硬件要求:由于涉及深度卷积计算,建议使用NVIDIA GPU。8GB以上显存可满足大多数标准模型推理和中小规模训练;若使用EfficientNet-B7、ConvNeXt等重型编码器,建议16GB+显存。CPU可运行推理,但训练速度会显著受限。
SMP本质上一个训练和推理框架,不提供开箱即用的标注工具或Web界面。你需要自行准备标注数据、编写训练循环、配置实验管理。此外,项目不支持实例分割(Instance Segmentation,如Mask R-CNN),仅限语义分割。对于需要实例级分割的任务,需要配合Detectron2或mmdetection使用。
SMP的出现极大地降低了图像语义分割的门槛。它证明了"统一接口 + 模块化组合"这一设计思路在深度学习工具库中的成功。12种分割架构与500+编码器的任意组合,使得研究者和工程师可以快速进行架构对比实验,而无需在代码适配上花费大量时间。
从发展趋势看,SMP正在向Transformer架构靠拢——Segformer、DPT等基于注意力机制的分割模型已纳入支持列表,与传统的CNN编码器形成互补。随着Vision Transformer的持续发展,SMP的编码器生态预计会进一步扩展。
对于希望快速将图像分割能力集成到产品中的开发者来说,SMP是目前PyTorch生态中最值得考虑的方案之一——成熟、稳定、文档完善、社区活跃。