首页>编程相关>其他源码
下载


具有图像神经网络的 Python 库
基于PyTorch的分割。



该库的主要特点是：

  • 高级 API（只需两行即可创建神经网络）
  • 10 个用于二元和多类分割的模型架构（包括传奇的 Unet）
  • 124 个可用编码器（以及 timm 的 500 多个编码器）
  • 所有编码器都有预先训练的权重，以实现更快更好的收敛
  • 训练例程的流行指标和损失

项目文档

访问阅读文档项目页面或阅读以下 README 以了解有关分割模型 Pytorch（简称 SMP）库的更多信息

？目录

  1. 快速启动
  2. 示例
  3. 型号
    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" ,        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights = "imagenet" ,     # use `imagenet` pre-trained weights for encoder initialization
    in_channels = 1 ,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes = 3 ,                      # model output channels (number of classes in your dataset)
)
  • 请参阅可用模型架构表
  • 请参阅可用编码器及其相应权重表

2. 配置数据预处理

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

 from segmentation_models_pytorch . encoders import get_preprocessing_fn

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

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

示例

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

？型号

架构

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

编码器

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

残差网络
编码器重量帕拉姆斯，M
资源网18图像网 / ssl / swsl 11M
资源网34图像网21M
resnet50图像网 / ssl / swsl 23M
资源网101图像网42M
资源网152图像网58M
剩余下一个
编码器重量帕拉姆斯，M
resnext50_32x4d图像网 / ssl / swsl 22M
resnext101_32x4d SSL/SWSL 42M
resnext101_32x8d imagenet/instagram/ssl/swsl 86M
resnext101_32x16d Instagram / ssl / swsl 191M
resnext101_32x32d Instagram 466M
resnext101_32x48d Instagram 826M
剩余巢穴
编码器重量帕拉姆斯，M
蒂姆·雷斯内斯特14d图像网8M
蒂姆-resnest26d图像网15M
蒂姆-雷斯内斯特50d图像网25M
蒂姆-resnest101e图像网46M
蒂姆-resnest200e图像网68M
蒂姆-resnest269e图像网108M
蒂姆-resnest50d_4s2x40d图像网28M
蒂姆-resnest50d_1s4x24d图像网23M
Res2Ne(X)t
编码器重量帕拉姆斯，M
timm-res2net50_26w_4s图像网23M
timm-res2net101_26w_4s图像网43M
timm-res2net50_26w_6s图像网35M
timm-res2net50_26w_8s图像网46M
timm-res2net50_48w_2s图像网23M
timm-res2net50_14w_8s图像网23M
timm-res2next50图像网22M
RegNet(x/y)
编码器重量帕拉姆斯，M
蒂姆-regnetx_002图像网2M
蒂姆-regnetx_004图像网4M
蒂姆-regnetx_006图像网5M
蒂姆-regnetx_008图像网6M
蒂姆-regnetx_016图像网8M
蒂姆-regnetx_032图像网14M
蒂姆-regnetx_040图像网20M
蒂姆-regnetx_064图像网24M
蒂姆-regnetx_080图像网37M
蒂姆-regnetx_120图像网43M
蒂姆-regnetx_160图像网52M
蒂姆-regnetx_320图像网105M
蒂姆-regnety_002图像网2M
蒂姆-regnety_004图像网3M
蒂姆-regnety_006图像网5M
蒂姆-regnety_008图像网5M
蒂姆-regnety_016图像网10M
蒂姆-regnety_032图像网17M
蒂姆-regnety_040图像网19M
蒂姆-regnety_064图像网29M
蒂姆-regnety_080图像网37M
蒂姆-regnety_120图像网49M
蒂姆-regnety_160图像网80M
蒂姆-regnety_320图像网141M
通用电气网络
编码器重量帕拉姆斯，M
蒂姆-gernet_s图像网6M
蒂姆-gernet_m图像网18M
蒂姆-gernet_l图像网28M
SE网络
编码器重量帕拉姆斯，M
塞内特154图像网113M
se_resnet50图像网26M
se_resnet101图像网47M
se_resnet152图像网64M
se_resnext50_32x4d图像网25M
se_resnext101_32x4d图像网46M
SK-ResNe(X)t
编码器重量帕拉姆斯，M
蒂姆-skresnet18图像网11M
蒂姆-skresnet34图像网21M
蒂姆-skresnext50_32x4d图像网25M
密集网络
编码器重量帕拉姆斯，M
密集网121图像网6M
密集网169图像网12M
密集网201图像网18M
密集网161图像网26M
盗梦空间
编码器重量帕拉姆斯，M
盗梦空间resnetv2 imagenet / imagenet+背景54M
盗梦空间v4 imagenet / imagenet+背景41M
异象图像网22M
高效网络
编码器重量帕拉姆斯，M
高效网-b0图像网4M
高效网-b1图像网6M
高效网-b2图像网7M
高效网-b3图像网10M
高效网络-b4图像网17M
高效网-b5图像网28M
高效网-b6图像网40M
高效网-b7图像网63M
timm-efficientnet-b0 imagenet / advprop / 嘈杂的学生4M
timm-efficientnet-b1 imagenet / advprop / 嘈杂的学生6M
timm-efficientnet-b2 imagenet / advprop / 嘈杂的学生7M
timm-efficientnet-b3 imagenet / advprop / 嘈杂的学生10M
timm-efficientnet-b4 imagenet / advprop / 嘈杂的学生17M
timm-efficientnet-b5 imagenet / advprop / 嘈杂的学生28M
timm-efficientnet-b6 imagenet / advprop / 嘈杂的学生40M
timm-efficientnet-b7 imagenet / advprop / 嘈杂的学生63M
timm-efficientnet-b8图像网 / advprop 84M
timm-efficientnet-l2吵闹的学生474M
timm-efficientnet-lite0图像网4M
timm-efficientnet-lite1图像网5M
timm-efficientnet-lite2图像网6M
timm-efficientnet-lite3图像网8M
timm-efficientnet-lite4图像网13M
移动网络
编码器重量帕拉姆斯，M
移动网络_v2图像网2M
timm-mobilenetv3_large_075图像网1.78M
timm-mobilenetv3_large_100图像网2.97M
timm-mobilenetv3_large_minimal_100图像网1.41M
timm-mobilenetv3_small_075图像网0.57M
timm-mobilenetv3_small_100图像网0.93M
timm-mobilenetv3_small_minimal_100图像网0.43M
DPN
编码器重量帕拉姆斯，M
dpn68图像网11M
dpn68b图像网+5k 11M
DPN92图像网+5k 34M
dpn98图像网58M
DPN107图像网+5k 84M
DPN131图像网76M
VGG
编码器重量帕拉姆斯，M
VG11图像网9M
vgg11_bn图像网9M
VG13图像网9M
vgg13_bn图像网9M
VG16图像网14M
vgg16_bn图像网14M
VG19图像网20M
vgg19_bn图像网20M
混合视觉变压器

来自 SegFormer 的骨干在 Imagenet 上进行了预训练！可以与包中的其他解码器一起使用，您可以将 Mix Vision Transformer 与 Unet、FPN 等结合使用！

限制：

  • Linknet、Unet++不支持编码器
  • FPN 仅支持编码器深度 = 5 的编码器
编码器重量帕拉姆斯，M
mit_b0图像网3M
mit_b1图像网13M
mit_b2图像网24M
mit_b3图像网44M
mit_b4图像网60M
mit_b5图像网81M
移动一号

Apple 的“亚一毫秒”Backbone 在 Imagenet 上进行了预训练！可与所有解码器一起使用。

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

编码器重量帕拉姆斯，M
mobileone_s0图像网4.6M
mobileone_s1图像网4.0M
mobileone_s2图像网6.5M
mobileone_s3图像网8.8M
mobileone_s4图像网13.6M

* ssl 、 swsl - ImageNet 上的半监督和弱监督学习（repo）。

蒂姆编码器

文档

Pytorch 图像模型（又名 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' ,             # one of 'avg', 'max'
    dropout = 0.5 ,               # dropout ratio, default is None
    activation = 'sigmoid' ,      # activation function, default is None
    classes = 4 ,                 # define number of output labels
)
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  # create .venv, install SMP in dev mode

运行测试和代码检查

make fixup         # Ruff for formatting and lint checks

使用编码器更新表

make table        # generate a table with encoders and print to stdout

引用

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

展开
附加信息
相关应用
为您推荐
相关资讯 全部