一大波PyTorch图像分割模型来袭，俄罗斯程序员出品新model zoo
鱼羊 发自 凹非寺
量子位 报道 | 公众号 QbitAI
一个新的图像分割model zoo来啦！
一大波基于PyTorch的图像分割模型整理好了就等你来用~
这个新集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造，包含四种模型架构和30种预训练骨干模型（backbone），官方文档列举了四条主要特点：
- 高级API（两行代码构建神经网络）
- 4种模型架构：Unet，Linknet，FPN，PSPNet
- 每个架构有30种可用的编码器
- 所有编码器都具有预先训练的权重，可实现更快、更好的收敛
授人以鱼不如授人以渔，放出了model zoo还不算，小哥还贴心地提供了如何使用CamVid数据集进行训练的示例。
效果怎么样
CamVid数据集是计算机视觉领域常用的一个数据集，通常用来进行街景分割。
俄罗斯小哥用这个数据集训练了一个能识别出图片中汽车的新模型。无论这辆车是隐藏在暗夜之中还是只有黄豆大小，都逃不出图像分割模型的手掌心。
构建模型时他选择了传奇的Unet架构，并使用了se_resnext50_32x4d这个预训练骨干模型。
听上去好像很厉害的样子，但群众都相信眼见为实，不看看模型训练效果如何，心里难免有一丢丢怀疑。
嗯，看来不管是昏天黑地还是目标只有一丁点大，新模型的预测结果都挺不错。
再加大难度呢？
量子位表示眼已瞎，模型是真强。
Pavel Yakubovskiy小哥还在Colab上放了用Catalyst（PyTorch上的高级框架）训练模型的代码，感兴趣的话可以直接在浏览器上跑一下。
入园指南
model zoo已开源，免票入场，快来围观。
先来了解一下都有哪些预训练模型可以用：
以及它们的权重：
pip一下：
$ pip install segmentation-models-pytorch
$ pip install git+https://github.com/qubvel/segmentation_models.pytorch
python库安装完成，就可以轻松地创建模型啦：
1import segmentation_models_pytorch as smp 2model = smp.Unet()
根据训练任务的不同，可以通过调整骨干模型来改变网络结构，并且使用预训练权重来进行初始化：
1model = smp.Unet('resnet34', encoder_weights='imagenet')
也可以改变模型的输出类型：
1model = smp.Unet('resnet34', classes=3, activation='softmax')
所有的模型都有预训练编码器，因此准备数据的方法必须与进行权重训练时一致：
1from segmentation_models_pytorch.encoders import get_preprocessing_fn 2preprocess_input = get_preprocessing_fn(‘renset18’, pretrained=’imagenet’)
One More Thing
其实Pavel Yakubovskiy小哥早在去年就发布了一个基于Keras的图像分割model zoo，一经推出就颇受欢迎。
一样的配方，一样的味道，简单易用，一脉相承。麻麻再也不用担心我找不到合适的预训练模型。
不过基于Keras的版本预训练模型要少一些，小哥也没有放出对应的示例，效果如何，就需要自己去检验啦~
传送门
PyTorch版Github地址：
https://github.com/qubvel/segmentation_models.pytorch
Keras版Github地址：
https://github.com/qubvel/segmentation_models
Colab上手：
https://colab.research.google.com/gist/Scitator/e3fd90eec05162e16b476de832500576/cars-segmentation-camvid.ipynb#scrollTo=A40deTWg-WJ1
- 0元起步打造你的AI搜索！实测秘塔知识库新功能，竟能指导我升职加薪了2024-11-11
- o1满血版泄露！奥数题图片推理手拿把掐，奥特曼上线剧透o22024-11-03
- 科研版AI搜索来了！知乎直答接入正版论文库，一手实测在此2024-11-01
- 微软清华改进Transformer：用降噪耳机原理升级注意力，一作在线答疑2024-11-03
相关阅读
