segmentation_models.pytorch

该库的主要特点包括：

• 高级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. 许可证

⏳ 快速开始

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')

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

💡 示例

• 使用Pytorch-Lightning训练宠物二分类分割模型的notebook和
• 在CamVid数据集上训练汽车分割模型在此。
• 使用Catalyst（PyTorch的高级框架）、TTAch（PyTorch的TTA库）和Albumentations（快速图像增强库）训练SMP模型 - 在此
• 使用Pytorch-Lightning框架训练SMP模型 - 在此（由@ternaus提供的服装二分类分割）。

📦 模型

架构

• Unet [论文] [文档]
• Unet++ [论文] [文档]
• MAnet [论文] [文档]
• Linknet [论文] [文档]
• FPN [论文] [文档]
• PSPNet [论文] [文档]
• PAN [论文] [文档]
• DeepLabV3 [论文] [文档]
• DeepLabV3+ [论文] [文档]

编码器

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

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

Timm编码器

文档

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

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

支持的编码器总数：549

• 可用编码器表格

🔁 模型API

• 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)

🛠 安装

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}}
}

🛡️ 许可证

项目在 MIT 许可证 下分发