该库的主要特性包括：

• 超级简单的高级 API（只需两行代码即可创建神经网络）
• 12 种编码器-解码器模型架构（Unet、Unet++、Segformer、DPT 等）
• 800 多个基于卷积和变换的 预训练 编码器，包括对 timm 的支持
• 用于训练过程的常用指标和损失函数（Dice、Jaccard、Tversky 等）
• 支持 ONNX 导出以及与 Torch Script/Trace/Compile 兼容

🤝 赞助商：withoutBG

withoutBG 是一款高质量的背景移除工具。他们使用 smp.Unet 构建了开源的图像抠图和精修模型，并自豪地赞助了这个项目。

📚 项目文档 📚

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

📋 目录

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

⏳ 快速入门

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

分割模型只是一个 PyTorch 的 torch.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，RGB 图像为 3 等）
    classes=3,                      // 模型输出通道数（数据集中类别数量）
)

• 请参阅 表格 了解可用的模型架构
• 请参阅 表格 了解可用的编码器及其对应的权重

2. 配置数据预处理

所有编码器都带有预训练权重。如果你的数据预处理方式与权重预训练时一致，可能会获得更好的结果（更高的指标得分和更快的收敛速度）。不过，如果你训练的是整个模型，而不仅仅是解码器，则这并非必需。

from segmentation_models_pytorch.encoders import get_preprocessing_fn

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

恭喜你！大功告成！现在你可以用自己喜欢的框架来训练你的模型了！

💡 示例

名称	链接	Colab
训练 牛津宠物数据集上的二分类分割	Notebook
训练 CamVid 数据集上的汽车二分类分割	Notebook
训练 CamVid 数据集上的多类分割	Notebook
训练 @ternaus 的服装二分类分割	仓库
加载并推理 预训练的 Segformer	Notebook
加载并推理 预训练的 DPT	Notebook
加载并推理 预训练的 UPerNet	Notebook
保存和加载 模型到本地或 HuggingFace Hub	Notebook
导出 训练好的模型为 ONNX	Notebook

📦 模型和编码器

架构

架构	论文	文档	检查点
Unet	论文	文档
Unet++	论文	文档
MAnet	论文	文档
Linknet	论文	文档
FPN	论文	文档
PSPNet	论文	文档
PAN	论文	文档
DeepLabV3	论文	文档
DeepLabV3+	论文	文档
UPerNet	论文	文档	检查点
Segformer	论文	文档	检查点
DPT	论文	文档	检查点

编码器

该库为分割模型提供了广泛的预训练编码器（也称为骨干网络）。我们不是使用分类模型最后一层的特征，而是提取中间特征并将其输入到解码器中，用于分割任务。

所有编码器都带有预训练权重，这有助于在训练分割模型时实现更快、更稳定的收敛。

鉴于支持的编码器种类繁多，您可以根据具体用例选择最适合的编码器，例如：

• 轻量级编码器适用于低延迟应用或边缘设备上的实时推理（如 mobilenet/mobileone）。
• 高容量架构适用于涉及大量分割类别的复杂任务，能够提供更高的精度（如 convnext/swin/mit）。

通过选择合适的编码器，您可以平衡效率、性能和模型复杂度，以满足项目需求。

所有编码器及其对应的预训练权重均列于文档中：

• 表格 中包含原生移植的编码器
• 表格 中包含由 timm 支持的编码器

🔁 模型 API

输入通道数

input_channels 参数允许您创建一个可以处理任意通道数张量的模型。
如果您使用 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 参数指定编码器中的下采样操作次数，因此您可以通过设置较小的 depth 来使模型更轻量化。

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

🛠 安装

PyPI 版本：

$ pip install segmentation-models-pytorch

GitHub 上的最新版本：

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

🏆 使用该库赢得的比赛

Segmentation Models 包广泛应用于图像分割竞赛中。
此处 列出了相关比赛、获胜者姓名以及他们的解决方案链接。

🛠 使用 SMP 构建的项目

• withoutBG：一个开源背景移除工具，其图像抠图和精修模型使用了 smp.Unet。请查看 withoutBG Focus on HuggingFace。

🤝 贡献

1. 以开发模式安装 SMP

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

2. 运行测试和代码检查

make test          # 使用 pytest 运行测试套件
make fixup         # 使用 Ruff 进行格式化和 lint 检查

3. 更新表格（如果您添加了新的编码器）

make table        # 生成编码器表格并打印到 stdout

📝 引用

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

🛡️ 许可证

该项目主要采用 MIT 许可证 发布，但部分文件可能适用其他许可证。请仔细查阅 LICENSES 文件及各文件中的许可声明，尤其是在商业用途时。

segmentation_models.pytorch 快速上手指南

segmentation_models.pytorch (简称 SMP) 是一个基于 PyTorch 的图像语义分割工具库。它提供了极简的高层 API，支持 12 种主流分割架构（如 Unet, Unet++, Segformer, DPT 等）以及 800+ 种预训练编码器（Backbones），能帮助开发者快速构建高性能的分割模型。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 3.9 或更高
• PyTorch 版本: 1.9 或更高
• 前置依赖: 已正确安装 CUDA（如需 GPU 加速）及对应的 torch 和 torchvision

建议先访问 PyTorch 官网 根据你的环境安装基础深度学习框架。

安装步骤

你可以使用 pip 直接安装最新稳定版。国内用户推荐使用清华或阿里镜像源以加速下载。

使用官方源安装：

pip install segmentation-models-pytorch

使用国内镜像源安装（推荐）：

pip install segmentation-models-pytorch -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

SMP 的核心优势在于其简洁的 API，仅需几行代码即可实例化一个完整的分割模型。

1. 创建模型

以下示例展示如何创建一个基于 ResNet34 编码器的 Unet 模型，并加载 ImageNet 预训练权重：

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,                      # 输出类别数 (根据数据集的类别数量设定)
)

提示：支持的架构包括 Unet, Unet++, MAnet, Linknet, FPN, PSPNet, DeepLabV3, DeepLabV3+, PAN, UPerNet, Segformer, DPT。

2. 配置数据预处理

为了充分发挥预训练权重的效果，建议对输入数据进行与预训练时相同的预处理（归一化等）。这通常能带来更快的收敛速度和更高的指标分数。

from segmentation_models_pytorch.encoders import get_preprocessing_fn

# 获取对应编码器的预处理函数
preprocess_input = get_preprocessing_fn('resnet34', pretrained='imagenet')

# 使用示例 (假设 image 是 numpy 数组)
# image = preprocess_input(image)

完成上述两步后，model 就是一个标准的 PyTorch torch.nn.Module，你可以直接使用熟悉的 PyTorch 训练循环、损失函数（如 Dice Loss, Jaccard Loss）和数据加载器进行模型训练或推理。