懂AI
segmentation_models.pytorch

segmentation_models.pytorch

基于PyTorch的神经网络图像分割库

segmentation_models.pytorch 是一个基于 PyTorch 的图像分割库,提供9种分割模型架构和124种编码器。该库 API 简洁,支持预训练权重,并包含常用评估指标和损失函数。它适用于研究和实际应用中的各种图像分割任务,是图像分割领域的实用工具。

图像分割PyTorch神经网络预训练模型编码器Github开源项目
访问官网GitHub论文文档
<div align="center">

logo
基于PyTorch的图像分割神经网络Python库

Generic badge GitHub Workflow Status (branch) Read the Docs <br> PyPI PyPI - Downloads <br> PyTorch - Version Python - Version

</div>

该库的主要特点包括：

  • 高级API（只需两行代码即可创建神经网络）
  • 9种用于二分类和多分类分割的模型架构（包括著名的Unet）
  • 124种可用编码器（以及来自timm的500多种编码器）
  • 所有编码器都有预训练权重，以实现更快更好的收敛
  • 用于训练的常用指标和损失函数

📚 项目文档 📚

访问Read The Docs项目页面或阅读以下README，了解更多关于Segmentation Models Pytorch（简称SMP）库的信息

📋 目录

  1. 快速开始
  2. 示例
  3. 模型
    1. 架构
    2. 编码器
    3. Timm编码器
  4. 模型API
    1. 输入通道
    2. 辅助分类输出
    3. 深度
  5. 安装
  6. 使用该库获胜的比赛
  7. 贡献
  8. 引用
  9. 许可证

⏳ 快速开始 <a name="start"></a>

1. 使用SMP创建你的第一个分割模型

分割模型只是一个PyTorch nn.Module，可以像这样轻松创建：

import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # 选择编码器，例如 mobilenet_v2 或 efficientnet-b7
    encoder_weights="imagenet",     # 使用`imagenet`预训练权重进行编码器初始化
    in_channels=1,                  # 模型输入通道数（1用于灰度图像，3用于RGB等）
    classes=3,                      # 模型输出通道数（数据集中的类别数）
)
  • 查看表格了解可用的模型架构
  • 查看表格了解可用的编码器及其对应的权重

2. 配置数据预处理

所有编码器都有预训练权重。以与权重预训练期间相同的方式准备数据可能会给你带来更好的结果（更高的指标分数和更快的收敛）。如果你训练整个模型而不仅仅是解码器，这不是必需的。

from segmentation_models_pytorch.encoders import get_preprocessing_fn

preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')

恭喜！你已完成！现在你可以使用你喜欢的框架来训练你的模型了！

💡 示例 <a name="examples"></a>

📦 模型 <a name="models"></a>

架构 <a name="architectures"></a>

编码器 <a name="encoders"></a>

以下是SMP支持的编码器列表。选择适当的编码器系列并点击展开表格，选择特定的编码器及其预训练权重（encoder_name和encoder_weights参数）。

<details> <summary style="margin-left: 25px;">ResNet</summary> <div style="margin-left: 25px;">
编码器权重参数量, M
resnet18imagenet / ssl / swsl11M
resnet34imagenet21M
resnet50imagenet / ssl / swsl23M
resnet101imagenet42M
resnet152imagenet58M
</div> </details> <details> <summary style="margin-left: 25px;">ResNeXt</summary> <div style="margin-left: 25px;">
编码器权重参数量, M
resnext50_32x4dimagenet / ssl / swsl22M
resnext101_32x4dssl / swsl42M
resnext101_32x8dimagenet / instagram / ssl / swsl86M
resnext101_32x16dinstagram / ssl / swsl191M
resnext101_32x32dinstagram466M
resnext101_32x48dinstagram826M
</div> </details> <details> <summary style="margin-left: 25px;">ResNeSt</summary> <div style="margin-left: 25px;">
编码器权重参数量, M
timm-resnest14dimagenet8M
timm-resnest26dimagenet15M
timm-resnest50dimagenet25M
timm-resnest101eimagenet46M
timm-resnest200eimagenet68M
timm-resnest269eimagenet108M
timm-resnest50d_4s2x40dimagenet28M
timm-resnest50d_1s4x24dimagenet23M
</div> </details> <details> <summary style="margin-left: 25px;">Res2Ne(X)t</summary> <div style="margin-left: 25px;">
编码器权重参数量, M
timm-res2net50_26w_4simagenet23M
timm-res2net101_26w_4simagenet43M
timm-res2net50_26w_6simagenet35M
timm-res2net50_26w_8simagenet46M
timm-res2net50_48w_2simagenet23M
timm-res2net50_14w_8simagenet23M
timm-res2next50imagenet22M
</div> </details> <details> <summary style="margin-left: 25px;">RegNet(x/y)</summary> <div style="margin-left: 25px;"> |编码器 |权重 |参数量（百万） | |--------------------------------|:------------------------------:|:------------------------------:| |timm-regnetx_002 |imagenet |2M | |timm-regnetx_004 |imagenet |4M | |timm-regnetx_006 |imagenet |5M | |timm-regnetx_008 |imagenet |6M | |timm-regnetx_016 |imagenet |8M | |timm-regnetx_032 |imagenet |14M | |timm-regnetx_040 |imagenet |20M | |timm-regnetx_064 |imagenet |24M | |timm-regnetx_080 |imagenet |37M | |timm-regnetx_120 |imagenet |43M | |timm-regnetx_160 |imagenet |52M | |timm-regnetx_320 |imagenet |105M | |timm-regnety_002 |imagenet |2M | |timm-regnety_004 |imagenet |3M | |timm-regnety_006 |imagenet |5M | |timm-regnety_008 |imagenet |5M | |timm-regnety_016 |imagenet |10M | |timm-regnety_032 |imagenet |17M | |timm-regnety_040 |imagenet |19M | |timm-regnety_064 |imagenet |29M | |timm-regnety_080 |imagenet |37M | |timm-regnety_120 |imagenet |49M | |timm-regnety_160 |imagenet |80M | |timm-regnety_320 |imagenet |141M | </div> </details> <details> <summary style="margin-left: 25px;">GERNet</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
timm-gernet_simagenet6M
timm-gernet_mimagenet18M
timm-gernet_limagenet28M
</div> </details> <details> <summary style="margin-left: 25px;">SE-Net</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
senet154imagenet113M
se_resnet50imagenet26M
se_resnet101imagenet47M
se_resnet152imagenet64M
se_resnext50_32x4dimagenet25M
se_resnext101_32x4dimagenet46M
</div> </details> <details> <summary style="margin-left: 25px;">SK-ResNe(X)t</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
timm-skresnet18imagenet11M
timm-skresnet34imagenet21M
timm-skresnext50_32x4dimagenet25M
</div> </details> <details> <summary style="margin-left: 25px;">DenseNet</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
densenet121imagenet6M
densenet169imagenet12M
densenet201imagenet18M
densenet161imagenet26M
</div> </details> <details> <summary style="margin-left: 25px;">Inception</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
inceptionresnetv2imagenet / imagenet+background54M
inceptionv4imagenet / imagenet+background41M
xceptionimagenet22M
</div> </details> <details> <summary style="margin-left: 25px;">EfficientNet</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
efficientnet-b0imagenet4M
efficientnet-b1imagenet6M
efficientnet-b2imagenet7M
efficientnet-b3imagenet10M
efficientnet-b4imagenet17M
efficientnet-b5imagenet28M
efficientnet-b6imagenet40M
efficientnet-b7imagenet63M
timm-efficientnet-b0imagenet / advprop / noisy-student4M
timm-efficientnet-b1imagenet / advprop / noisy-student6M
timm-efficientnet-b2imagenet / advprop / noisy-student7M
timm-efficientnet-b3imagenet / advprop / noisy-student10M
timm-efficientnet-b4imagenet / advprop / noisy-student17M
timm-efficientnet-b5imagenet / advprop / noisy-student28M
timm-efficientnet-b6imagenet / advprop / noisy-student40M
timm-efficientnet-b7imagenet / advprop / noisy-student63M
timm-efficientnet-b8imagenet / advprop84M
timm-efficientnet-l2noisy-student474M
timm-efficientnet-lite0imagenet4M
timm-efficientnet-lite1imagenet5M
timm-efficientnet-lite2imagenet6M
timm-efficientnet-lite3imagenet8M
timm-efficientnet-lite4imagenet13M
</div> </details> <details> <summary style="margin-left: 25px;">MobileNet</summary> <div style="margin-left: 25px;"> |编码器 |权重 |参数量（百万） | |--------------------------------|:------------------------------:|:------------------------------:| |mobilenet_v2 |imagenet |2M | |timm-mobilenetv3_large_075 |imagenet |1.78M | |timm-mobilenetv3_large_100 |imagenet |2.97M | |timm-mobilenetv3_large_minimal_100|imagenet |1.41M | |timm-mobilenetv3_small_075 |imagenet |0.57M | |timm-mobilenetv3_small_100 |imagenet |0.93M | |timm-mobilenetv3_small_minimal_100|imagenet |0.43M | </div> </details> <details> <summary style="margin-left: 25px;">DPN</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
dpn68imagenet11M
dpn68bimagenet+5k11M
dpn92imagenet+5k34M
dpn98imagenet58M
dpn107imagenet+5k84M
dpn131imagenet76M
</div> </details> <details> <summary style="margin-left: 25px;">VGG</summary> <div style="margin-left: 25px;">
编码器权重参数量（百万）
vgg11imagenet9M
vgg11_bnimagenet9M
vgg13imagenet9M
vgg13_bnimagenet9M
vgg16imagenet14M
vgg16_bnimagenet14M
vgg19imagenet20M
vgg19_bnimagenet20M
</div> </details> <details> <summary style="margin-left: 25px;">混合视觉Transformer</summary> <div style="margin-left: 25px;">

SegFormer在Imagenet上预训练的骨干网络！可以与包中的其他解码器一起使用，您可以将混合视觉Transformer与Unet、FPN等结合使用！

限制：

  • 编码器不支持Linknet、Unet++
  • 编码器仅支持FPN，且仅限编码器深度为5的情况
编码器权重参数量（百万）
mit_b0imagenet3M
mit_b1imagenet13M
mit_b2imagenet24M
mit_b3imagenet44M
mit_b4imagenet60M
mit_b5imagenet81M
</div> </details> <details> <summary style="margin-left: 25px;">MobileOne</summary> <div style="margin-left: 25px;">

苹果公司的"亚毫秒级"骨干网络在Imagenet上预训练！可以与所有解码器一起使用。

注意：在官方GitHub仓库中，s0变体有额外的num_conv_branches，导致比s1有更多的参数。

编码器权重参数量（百万）
mobileone_s0imagenet4.6M
mobileone_s1imagenet4.0M
mobileone_s2imagenet6.5M
mobileone_s3imagenet8.8M
mobileone_s4imagenet13.6M
</div> </details>

* ssl、swsl - 在ImageNet上进行半监督和弱监督学习（仓库）。

Timm编码器 <a name="timm"></a>

文档

Pytorch Image Models（简称timm）有许多预训练模型和接口，可以将这些模型用作smp中的编码器，但并非所有模型都受支持

  • 并非所有transformer模型都实现了编码器所需的features_only功能
  • 一些模型的步幅不适合

支持的编码器总数：549

🔁 模型API <a name="api"></a>

  • model.encoder - 预训练的骨干网络，用于提取不同空间分辨率的特征
  • model.decoder - 取决于模型架构（Unet/Linknet/PSPNet/FPN）
  • model.segmentation_head - 最后一个块，用于生成所需数量的掩码通道（还包括可选的上采样和激活）
  • model.classification_head - 可选块，在编码器顶部创建分类头
  • model.forward(x) - 顺序地将x通过模型的编码器、解码器和分割头（如果指定，还包括分类头）
输入通道

输入通道参数允许您创建可处理任意数量通道张量的模型。如果您使用来自imagenet的预训练权重 - 第一个卷积的权重将被重用。对于1通道情况，它将是第一个卷积层权重的总和，否则通道将使用如下方式填充权重：new_weight[:, i] = pretrained_weight[:, i % 3]，然后用new_weight * 3 / new_in_channels进行缩放。

model = smp.FPN('resnet34', in_channels=1)
mask = model(torch.ones([1, 1, 64, 64]))
辅助分类输出

所有模型都支持aux_params参数，默认设置为None。如果aux_params = None，则不创建分类辅助输出，否则模型不仅产生mask，还产生形状为NC的label输出。分类头由GlobalPooling->Dropout(可选)->Linear->Activation(可选)层组成，可以通过aux_params进行配置，如下所示：

aux_params=dict(
    pooling='avg',             # 'avg'或'max'之一
    dropout=0.5,               # dropout比率，默认为None
    activation='sigmoid',      # 激活函数，默认为None
    classes=4,                 # 定义输出标签的数量
)
model = smp.Unet('resnet34', classes=4, aux_params=aux_params)
mask, label = model(x)
深度

深度参数指定编码器中的下采样操作数量，因此如果指定较小的depth，可以使模型更轻量。

model = smp.Unet('resnet34', encoder_depth=4)

🛠 安装 <a name="installation"></a>

PyPI 版本：

$ pip install segmentation-models-pytorch

从源代码安装最新版本：

$ pip install git+https://github.com/qubvel/segmentation_models.pytorch

🏆 使用本库获得的比赛胜利

Segmentation Models包在图像分割比赛中被广泛使用。 这里你可以找到比赛名称、获胜者姓名以及他们解决方案的链接。

🤝 贡献

安装 SMP

make install_dev  # 创建 .venv，以开发模式安装 SMP

运行测试和代码检查

make fixup         # 使用 Ruff 进行格式化和代码检查

更新编码器表格

make table        # 生成编码器表格并输出到标准输出

📝 引用

@misc{Iakubovskii:2019,
  Author = {Pavel Iakubovskii},
  Title = {Segmentation Models Pytorch},
  Year = {2019},
  Publisher = {GitHub},
  Journal = {GitHub repository},
  Howpublished = {\url{https://github.com/qubvel/segmentation_models.pytorch}}
}

🛡️ 许可证 <a name="license"></a>

项目在 MIT 许可证 下分发

编辑推荐精选

Keevx

Keevx

AI数字人视频创作平台

Keevx 一款开箱即用的AI数字人视频创作平台，广泛适用于电商广告、企业培训与社媒宣传，让全球企业与个人创作者无需拍摄剪辑，就能快速生成多语言、高质量的专业视频。

即梦AI

即梦AI

一站式AI创作平台

提供 AI 驱动的图片、视频生成及数字人等功能，助力创意创作

扣子-AI办公

扣子-AI办公

AI办公助手，复杂任务高效处理

AI办公助手，复杂任务高效处理。办公效率低？扣子空间AI助手支持播客生成、PPT制作、网页开发及报告写作，覆盖科研、商业、舆情等领域的专家Agent 7x24小时响应，生活工作无缝切换，提升50%效率！

TRAE编程

TRAE编程

AI辅助编程，代码自动修复

Trae是一种自适应的集成开发环境（IDE），通过自动化和多元协作改变开发流程。利用Trae，团队能够更快速、精确地编写和部署代码，从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能，是提升开发效率的理想工具。

AI工具TraeAI IDE协作生产力转型热门
蛙蛙写作

蛙蛙写作

AI小说写作助手，一站式润色、改写、扩写

蛙蛙写作—国内先进的AI写作平台，涵盖小说、学术、社交媒体等多场景。提供续写、改写、润色等功能，助力创作者高效优化写作流程。界面简洁,功能全面,适合各类写作者提升内容品质和工作效率。

AI辅助写作AI工具蛙蛙写作AI写作工具学术助手办公助手营销助手AI助手
问小白

问小白

全能AI智能助手，随时解答生活与工作的多样问题

问小白，由元石科技研发的AI智能助手，快速准确地解答各种生活和工作问题，包括但不限于搜索、规划和社交互动，帮助用户在日常生活中提高效率，轻松管理个人事务。

热门AI助手AI对话AI工具聊天机器人
Transly

Transly

实时语音翻译/同声传译工具

Transly是一个多场景的AI大语言模型驱动的同声传译、专业翻译助手，它拥有超精准的音频识别翻译能力，几乎零延迟的使用体验和支持多国语言可以让你带它走遍全球，无论你是留学生、商务人士、韩剧美剧爱好者，还是出国游玩、多国会议、跨国追星等等，都可以满足你所有需要同传的场景需求，线上线下通用，扫除语言障碍，让全世界的语言交流不再有国界。

讯飞智文

讯飞智文

一键生成PPT和Word，让学习生活更轻松

讯飞智文是一个利用 AI 技术的项目，能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定，还是学生群体的职业生涯规划、实习避坑指南，亦或是活动策划、旅游攻略等内容，它都能提供支持，帮助用户精准表达，轻松呈现各种信息。

AI办公办公工具AI工具讯飞智文AI在线生成PPTAI撰写助手多语种文档生成AI自动配图热门
讯飞星火

讯飞星火

深度推理能力全新升级，全面对标OpenAI o1

科大讯飞的星火大模型，支持语言理解、知识问答和文本创作等多功能，适用于多种文件和业务场景，提升办公和日常生活的效率。讯飞星火是一个提供丰富智能服务的平台，涵盖科技资讯、图像创作、写作辅助、编程解答、科研文献解读等功能，能为不同需求的用户提供便捷高效的帮助，助力用户轻松获取信息、解决问题，满足多样化使用场景。

热门AI开发模型训练AI工具讯飞星火大模型智能问答内容创作多语种支持智慧生活
Spark-TTS

Spark-TTS

一种基于大语言模型的高效单流解耦语音令牌文本到语音合成模型

Spark-TTS 是一个基于 PyTorch 的开源文本到语音合成项目，由多个知名机构联合参与。该项目提供了高效的 LLM（大语言模型）驱动的语音合成方案，支持语音克隆和语音创建功能，可通过命令行界面（CLI）和 Web UI 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数，生成高质量的语音。该项目适用于多种场景，如有声读物制作、智能语音助手开发等。

下拉加载更多