AugustMe的学习小课堂 2020-10-20 10:22:33 912 收藏 1 分类专栏: PyTorch 图像分割 版权
segmentation_models_pytorch是一个基于PyTorch的图像分割神经网络
这个新集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造。
github地址:https://github.com/qubvel/segmentation_models.pytorch
该库的主要功能是:
-
高级API(只需两行即可创建神经网络)
-
用于二分类和多类分割的7种模型架构(包括传奇的Unet)
-
每种架构有57种可用的编码器
-
所有编码器均具有预训练的权重,以实现更快更好的收敛
一、安装
PyPI version:
pip install segmentation-models-pytorch
- 1
Latest version from source:
pip install git+https://github.com/qubvel/segmentation_models.pytorch
- 1
二、使用
由于该库是基于PyTorch框架构建的,因此创建的细分模型只是一个PyTorch nn.Module,可以轻松地创建它:
import segmentation_models_pytorch as smp
model = smp.Unet()
- 1
- 2
根据任务的不同,您可以通过选择具有更少或更多参数的主干并使用预训练的权重来初始化它来更改网络体系结构:
model = smp.Unet('resnet34', encoder_weights='imagenet')
- 1
更改模型中输出类的数量:
model = smp.Unet('resnet34', classes=3, activation='softmax')
- 1
所有模型均具有预训练的编码器,因此您必须按照权重预训练的相同方法准备数据:
from segmentation_models_pytorch.encoders import get_preprocessing_fn
preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')
- 1
- 2
更加详细的请见github源码。
注,还有其他优秀的库
例如https://github.com/CSAILVision/semantic-segmentation-pytorch
真的是学无止境啊,先star吧,再继续学习。
参考:
https://baijiahao.baidu.com/s?id=1632513555001646060&wfr=spider&for=pc
https://github.com/qubvel/segmentation_models.pytorch
https://github.com/CSAILVision/semantic-segmentation-pytorch
https://zhuanlan.zhihu.com/p/80688663
(这里面有一个实战项目)