segmentation_models_pytorch是一个基于PyTorch的图像分割神经网络
这个新集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造。 github地址:https://github.com/qubvel/segmentation_models.pytorch
segmentation_models_pytorch主要功能是:
高级API(只需两行即可创建神经网络)
用于二分类和多类分割的7种模型架构(包括传奇的Unet)
每种架构有57种可用的编码器
所有编码器均具有预训练的权重,以实现更快更好的收敛
一、安装 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')
|