懂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 许可证 下分发

编辑推荐精选

Trae

Trae

字节跳动发布的AI编程神器IDE

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

AI工具TraeAI IDE协作生产力转型热门
问小白

问小白

全能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 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数，生成高质量的语音。该项目适用于多种场景，如有声读物制作、智能语音助手开发等。

咔片PPT

咔片PPT

AI助力，做PPT更简单！

咔片是一款轻量化在线演示设计工具，借助 AI 技术，实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT，提供海量模板、智能美化、素材替换等功能，适用于销售、教师、学生等各类人群，能高效制作出高品质 PPT，满足不同场景演示需求。

讯飞绘文

讯飞绘文

选题、配图、成文，一站式创作，让内容运营更高效

讯飞绘文，一个AI集成平台，支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容，加速品牌传播，提升内容营销效果。

热门AI辅助写作AI工具讯飞绘文内容运营AI创作个性化文章多平台分发AI助手
材料星

材料星

专业的AI公文写作平台，公文写作神器

AI 材料星，专业的 AI 公文写作辅助平台，为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能，支持 30 + 文稿类型生成，助力快速完成领导讲话、工作总结、述职报告等材料，提升办公效率，是体制打工人的得力写作神器。

openai-agents-python

openai-agents-python

OpenAI Agents SDK，助力开发者便捷使用 OpenAI 相关功能。

openai-agents-python 是 OpenAI 推出的一款强大 Python SDK，它为开发者提供了与 OpenAI 模型交互的高效工具，支持工具调用、结果处理、追踪等功能，涵盖多种应用场景，如研究助手、财务研究等，能显著提升开发效率，让开发者更轻松地利用 OpenAI 的技术优势。

下拉加载更多