直接使用
请打开基于EasyCV的Mask2Former图像分割示例,并点击右上角 “ 在DSW中打开” 。
EasyCV图像分割-Mask2Former
Mask2Former是一个统一的Transformer-based图像分割架构,能够在相同的模型结构下实现语义分割、实例分割以及全景分割任务。
本文将介绍如何利用EasyCV基于Mask2Former进行图像分割模型的训练,以及如何利用训练好的模型进行图像分割预测
运行环境要求
PAI-Pytorch镜像 or 原生Pytorch1.8+以上环境 GPU机器, 内存32G以上
安装依赖包
注: 在PAI-DSW docker中无需安装相关依赖,可跳过此部分 在本地notebook环境中执行
1、 首先,安装pytorch和对应版本的torchvision,支持Pytorch1.8以上版本
# install pytorch and torch vision ! conda install --yes pytorch==1.10.0 torchvision==0.11.0 -c pytorch
2、获取torch和cuda版本,安装对应版本的mmcv和nvidia-dali
import torch import os os.environ['CUDA']='cu' + torch.version.cuda.replace('.', '') os.environ['Torch']='torch'+torch.version.__version__.replace('+PAI', '') !echo "cuda version: $CUDA" !echo "pytorch version: $Torch"
# install some python deps ! pip install mmdet ! pip install mmcv-full==1.4.4 -f https://download.openmmlab.com/mmcv/dist/${CUDA}/${Torch}/index.html ! pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/third_party/nvidia_dali_cuda100-0.25.0-1535750-py3-none-manylinux2014_x86_64.whl
3、 安装EasyCV算法包
pip install pai-easycv
4、 简单验证
from easycv.apis import *
5、安装deformable_attention
import easycv print(easycv.__file__) # 进入easycv安装目录编译deformable_attention ! cd /home/pai/lib/python3.6/site-packages/easycv/thirdparty/deformable_attention && python setup.py build install
数据准备
接下来介绍基于coco数据集的实例分割训练示例,你可以下载COCO2017数据,也可以使用我们提供了示例COCO数据
! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz
# 重命名文件夹 ! mkdir -p data/ && mv small_coco_demo database/coco
data/coco格式如下
data/coco/ ├── annotations │ ├── instances_train2017.json │ └── instances_val2017.json ├── train2017 │ ├── 000000005802.jpg │ ├── 000000060623.jpg │ ├── 000000086408.jpg │ ├── 000000118113.jpg │ ├── 000000184613.jpg │ ├── 000000193271.jpg │ ├── 000000222564.jpg │ ... │ └── 000000574769.jpg └── val2017 ├── 000000006818.jpg ├── 000000017627.jpg ├── 000000037777.jpg ├── 000000087038.jpg ├── 000000174482.jpg ├── 000000181666.jpg ├── 000000184791.jpg ├── 000000252219.jpg ... └── 000000522713.jpg
模型训练
这个Demo中我们采用Mask2Former图像分割算法训练ResNet50主干网络, 下载示例配置文件
! rm -rf mask2former_r50_8xb2_e50_instance.py ! wget https://raw.githubusercontent.com/alibaba/EasyCV/master/configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py
为了适配小数据,我们对配置文件mask2former_r50_8xb2_e50_instance.py做如下字段的修改,减少训练epoch数目,加大打印日志的频率
total_epochs = 3 #optimizer.lr -> 0.000001 # optimizer optimizer = dict( type='AdamW', lr=0.000001, weight_decay=0.05, eps=1e-8, betas=(0.9, 0.999), paramwise_options={ 'backbone': dict(lr_mult=0.1), 'query_embed': dict(weight_decay=0.), 'query_feat': dict(weight_decay=0.), 'level_embed': dict(weight_decay=0.), 'norm': dict(weight_decay=0.), }) # log_config.interval 1 log_config = dict(interval=1)
注意: 如果是使用COCO完整数据训练,为了保证效果,建议使用单机8卡进行训练;
为了保证模型效果,我们在预训练模型基础上finetune, 执行如下命令启动训练
!python -m easycv.tools.train mask2former_r50_8xb2_e50_instance.py --work_dir work_dir/segmentatino/mask2former_r50_instance --load_from http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_instance/epoch_50.pth
预测
下载测试图片
! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/000000123213.jpg
使用训练好的模型进行图像分割预测
import torch import cv2 from IPython.display import Image from easycv.predictors.segmentation import Mask2formerPredictor predictor = Mask2formerPredictor(model_path='work_dir/segmentatino/mask2former_r50_instance/epoch_3.pth', config_file='mask2former_r50_8xb2_e50_instance.py', task_mode='instance') img = cv2.imread('000000123213.jpg') predict_out = predictor(['000000123213.jpg']) instance_img = predictor.show_instance(img, **predict_out[0]) cv2.imwrite('instance_out.jpg',instance_img) display(Image('000000123213.jpg')) display(Image('instance_out.jpg'))