图像分割库segmentation_models.pytorch

segmentation_models_pytorch是一个基于PyTorch的图像分割神经网络

这个新集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造。
github地址：https://github.com/qubvel/segmentation_models.pytorch 

该库的主要功能是：

1. 高级API（只需两行即可创建神经网络）

2. 用于二分类和多类分割的9种模型架构（包括传奇的Unet）

3. 每种架构有104种可用的编码器

4. 所有编码器均具有预训练的权重，以实现更快更好的收敛

一、安装

PyPI version:

pip install segmentation-models-pytorch

 

• Latest version from source:

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

 

二、使用

• 由于该库是基于PyTorch框架构建的，因此创建的细分模型只是一个PyTorch nn.Module，可以轻松地创建它：

  

   

    

     

    
    

     

      import segmentation_models_pytorch as smp
     
    
   

    

     

    
    

     

      model = smp.Unet()
     
    
  
 

• 根据任务的不同，您可以通过选择具有更少或更多参数的主干并使用预训练的权重来初始化它来更改网络体系结构：

model = smp.Unet('resnet34', encoder_weights='imagenet')

 

• 更改模型中输出类的数量：

model = smp.Unet('resnet34', classes=3, activation='softmax')

 

• 所有模型均具有预训练的编码器，因此您必须按照权重预训练的相同方法准备数据：

  

   

    

     

    
    

     

      from segmentation_models_pytorch.encoders import get_preprocessing_fn
     
    
   

    

     

    
    

     

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

编码器列表

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

ResNet

ResNeXt

ResNeSt

Res2Ne(X)t

RegNet(x/y)

SE-Net

SK-ResNe(X)t

DenseNet

Inception

EfficientNet

MobileNet

DPN

VGG

文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。

原文链接：wanghao.blog.csdn.net/article/details/116245088