首頁>編程相關>其他源碼
下載


具有影像神經網路的 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
蒂姆-resnest14d影像網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 許可證分發

展開
附加信息
相關應用
爲您推薦
相關資訊 全部