最近涉及到了图像分割的任务,于是拿来写下博客加深下使用。
MMsegmentation与MMdeploy的环境配置暂不做讲解,在官网和其他博客中有很多说明。
MMdeploy主要是把pt转为 onnx_int8的情况。
MMsegmentation环境配置可以参考 : 安装与配置MMSegmentation
目录
- MMsegmentation简易使用
- 数据集的准备
- 数据集的config设置
- 新建一个自己的dataset类
- 修改__init__.py
- 新建自己config文件
- model文件设置
- 训练
- 测试
- MMdeploy简易使用
- deploy_cfg
- deploy.py
- 问题及改进
- 参考博客
- 欢迎指正
MMsegmentation简易使用
数据集的准备
采用MICCAI2023牙齿分割挑战的数据集----CHASE_DB1
以下的例子都是针对这个数据集,算是一个二分类的图像分割任务(背景,牙齿)
CHASE_DB1_牙齿分割数据集
我将其放到了如下位置:
work_dir/tt/mmsegmentation-main/datasets/CHASE_DB1/
数据集的config设置
新建一个自己的dataset类
路径如下 : work_dir/mmsegmentation-main/mmseg/datasets
新建一个tooth_test_dataset.py
可以复制同目录下的 cityscapes.py,然后在它上面进行如下3处修改
修改__init__.py
work_dir/mmsegmentation-main/mmseg/datasets/init.py
在如下两个位置加上刚才的类名,来供后续的代码调用。
新建自己config文件
work_dir/mmsegmentation-main/configs/base/datasets/
下新建自己的config文件 tooth_test_640x320.py (结尾可以是W x H的形式,方便后续更改)
可以复制同目录下的cityscapes.py进行更改
更改:dataset_type; data_root ; crop_size
train_dataloader中 Img_path 和 seg_map_path路径的设置
其中batch_size 和num_workers可以自行设置
model文件设置
上面这样就算是把数据集相关的文件设置好了,现在来配置Model文件的设置
work_dir/mmsegmentation-main/configs
下面有很多种模型,读者可以自行选择一个,笔者选择的是hrnet模型
新建如下py:work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-40k_tooth_test-640x320.py
可以复制同目录下的fcn_hr18_4xb2-40k_cityscapes-512x1024.py 然后进行修改
同样的还有crop_size也需要修改,要和数据集中的大小设置为一致。
然后来看__base__第一个文件,fcn_hr18_cls2.py
新建/home/siyingzhen/tt/mmsegmentation-main/configs/base/models/fcn_hr18_cls2.py
读者可以通过复制同目录下fcn_hr18.py进行修改
修改对应类别即可
然后来看__base__第二个文件,这是第一步 数据集的config设置的文件,没什么好讲的
然后来看__base__第四个文件,这是训练时的一些超参数:
训练
使用.tools/train.py或者.tools/dist_train.sh进行训练
接下来就可以进行训练了
单gpu训练:
mmsegmentation-main]$ python tools/train.py work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-20k_tooth_test-640x320.py --work-dir work_dir
B U G ,笔者不会解决 \textcolor{red}{BUG,笔者不会解决} BUG,笔者不会解决 :
如果使用–resume会使得刚开始加载数据的时候卡住
官方ISSUE
多GPU训练:
CUDA_VISIBLE_DEVICES=3,4,5,6 ./tools/dist_train.sh work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-20k_tooth_test-640x320.py 4
测试
使用.tools/test.py进行测试
python tools/test.py work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-40k_tooth_test-640x320.py work_dir/mmsegmentation-main/work_dirs/fcn_hr18_4xb2-40k_tooth_test-640x320/iter_20000.pth --show-dir work_dir/mmsegmentation-main/work_dirs/eval_show
然后就可以看到测试结果可视化情况
MMdeploy简易使用
环境配置相关还是不说了
参考:mmdeploy环境配置
(笔者是参考的官方github)
目标,mmsegmentation的pt模型转为onnx_int8
主要代码:
work_dir/deploy-main/tools/deploy.py
需要的是4个参数
deploy_cfg
/home/siyingzhen/tt/mmdeploy-main/configs/mmseg/segmentation_onnxruntime-int8_dynamic.py
需要修改__base__的第二个文件路径
需要修改input_shape中,一般是训练时的图片大小(640x320)
.
./segmentation_static.py如下:
再往下看…/base/onx_config.py
会发现这里和
torch.onnx.export的API参数设置基本一致,实际上deploy.py中如果选择转ONNX,最终执行的就是这个接口
现在回过头来看…/base/backends/onnxruntime_int8.py
work_dir/mmdeploy-main/configs/base/backends/onnxruntime-int8.py
笔者是复制了同目录下的onnxruntime-fp16.py然后进行修改的
deploy.py
于此,就得到了deploy_cfg。
model_cfg在mmsegmentation中我们已经设置了,在work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-40k_tooth_test-640x320.py
checkpoint就是训练的模型权重结果
img 可以是训练中的任意一张图片
使用代码:
python tools/deploy.py
work_dir/mmdeploy-main/configs/mmseg/segmentation_onnxruntime_static-640x320.py
work_dir/mmsegmentation-main/configs/hrnet/fcn_hr18_4xb2-40k_tooth_test-640x320.py
work_dir/mmsegmentation-main/work_dirs/fcn_hr18_4xb2-40k_tooth_test-640x320/iter_20000.pth
work_dir/mmsegmentation-main/datasets/CHASE_DB1/images/training/A-1.png
--work-dir work_dir/mmdeploy-main/work_dir/HRnet_640x320
最终可能会报错(主要是可视化方面的报错),但是也会生成量化后的onnx
Error:
RuntimeError: Exporting the operator einsum to ONNX opset version 11 is not supported. Support for this operator was added in version 12, try exporting with this version.
问题及改进
最后发现,这个onnx里面的数据还是fp32的。
如果进行如下修改:
会发现onnx中的数据是fp16的,这个是成功的。
后来问了下别人,如果要从pt量化为onnx int8的话最好使用 onnxruntime.quantization 下 quantize_static或quantize_dynamic的函数
以下是另一个大佬写的mmdeploy环境下,将pt转为onnx_int8的代码,读者可以参考
(涉及了Onnxsim简化 以及 核心的quantize_static)
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import logging
import os
import os.path as osp
import numpy as np
import random
import subprocess
import sys
from mmdeploy.apis import (extract_model, get_predefined_partition_cfg,visualize_model,torch2onnx)
from mmdeploy.utils import get_input_shape
from mmdeploy.utils import (get_ir_config, get_backend_config,get_common_config,Backend,get_root_logger, load_config)from mmengine.runner import Runner
from mmengine.registry import DATASETS
from onnxruntime.quantization import CalibrationDataReader, QuantFormat, quantize_static, QuantType, CalibrationMethod
from mmdeploy.apis.utils import build_task_processordef batch_reader(datas,task_processor,data_preprocessor,input_shape,batch_size):_datas = []length = len(datas)# print("len: ",length)# exit()max_num = 20 #建议设置到300-500之间,此处设置为20是为了加速。print("max_num: ",max_num)for i, data in enumerate(datas):if i>max_num:return Noneimg_path = data['data_samples'].img_pathdata, model_inputs = task_processor.create_input(img_path,input_shape,data_preprocessor=data_preprocessor)if isinstance(model_inputs, list) and len(model_inputs) == 1:model_inputs = model_inputs[0]if batch_size==1:yield {'input': model_inputs.numpy()}elif (i+1) % batch_size==0:_datas.append(data)yield {'input': np.concatenate(_datas, 0)}_datas = []elif i<length-1:_datas.append(model_inputs.numpy())else:_datas.append(model_inputs.numpy())yield {'input': np.concatenate(_datas, 0)}class DataReader(CalibrationDataReader):def __init__(self, dataset,task_processor,data_preprocessor,input_shape,batch_size=1):self.datas = batch_reader(dataset,task_processor,data_preprocessor,input_shape, batch_size)def get_next(self):return next(self.datas, None)def parse_args():parser = argparse.ArgumentParser(description='Export model to ONNX.')# parser.add_argument('--deploy_cfg',default='/root/mmdeploy/configs/mmseg/segmentation_onnxruntime_int8_static-512x512.py', help='deploy config path')# parser.add_argument('--model_cfg',default='model_weights/fcn/fcn_r50-d8_4xb4-20k_voc12aug-512x512.py',help='model config path')# parser.add_argument('--checkpoint',default='model_weights/fcn/fcn_r50-d8_512x512_20k_voc12aug_20200617_010715-52dc5306.pth',help='model checkpoint path')# parser.add_argument('--test_img',default='/root/mmdeploy/data/VOCdevkit/VOC2012/JPEGImages/2007_000027.jpg',help='image used to convert model model')parser.add_argument('--deploy_cfg',default='', help='deploy config path')parser.add_argument('--model_cfg',default='',help='model config path')parser.add_argument('--checkpoint',default='',help='model checkpoint path')parser.add_argument('--test_img',default='',help='image used to convert model model')parser.add_argument('--work-dir',default='',help='Directory to save output files.')parser.add_argument('--device', help='device used for conversion', default='cpu')parser.add_argument('--log-level',help='set log level',default='INFO',choices=list(logging._nameToLevel.keys()))args = parser.parse_args()return argsdef main():args = parse_args()logger = get_root_logger(log_level=args.log_level)logger.info(f'torch2onnx: \n\tmodel_cfg: {args.model_cfg} 'f'\n\tdeploy_cfg: {args.deploy_cfg}')os.makedirs(args.work_dir, exist_ok=True)# load deploy_cfgdeploy_cfg, model_cfg = load_config(args.deploy_cfg, args.model_cfg)save_file = get_ir_config(deploy_cfg)['save_file']torch2onnx(args.test_img,args.work_dir,save_file,deploy_cfg=args.deploy_cfg,model_cfg=args.model_cfg,model_checkpoint=args.checkpoint,device=args.device)backend_cfg = get_backend_config(deploy_cfg)precision = backend_cfg.get('precision', 'fp32')if precision == 'fp16':import onnxfrom onnxconverter_common import float16common_cfg = get_common_config(deploy_cfg)model = onnx.load(os.path.join(args.work_dir,save_file))model_fp16 = float16.convert_float_to_float16(model, **common_cfg)onnx.save(model_fp16, os.path.join(args.work_dir,save_file.replace('.onnx','_fp16.onnx')))if precision == 'int8':dataset = DATASETS.build(model_cfg['train_dataloader']['dataset'])task_processor = build_task_processor(model_cfg, deploy_cfg, 'cpu')torch_model = task_processor.build_pytorch_model(args.checkpoint)data_preprocessor=getattr(torch_model, 'data_preprocessor', None)input_shape = get_input_shape(deploy_cfg)print("input_shape: ",input_shape)data_reader = DataReader(dataset,task_processor,data_preprocessor,input_shape,1)sim_onnx_path = os.path.join(args.work_dir,save_file.replace('.onnx','_sim.onnx'))if os.path.exists(os.path.join(args.work_dir,save_file)):command = r"/root/miniconda3/envs/mmdeploy/bin/python -m onnxsim {} {}".format(os.path.join(args.work_dir,save_file),sim_onnx_path)process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)for line in process.stdout:print(line.rstrip())for line in process.stderr:print(line.rstrip(), file=sys.stderr)process.wait()quant_onnx_path = sim_onnx_path.replace("_sim.onnx","_sim_int8.onnx")if os.path.join(sim_onnx_path):quantize_static(model_input=sim_onnx_path, # 输入模型model_output=quant_onnx_path, # 输出模型calibration_data_reader=data_reader, # 校准数据读取器quant_format= QuantFormat.QDQ, # 量化格式 QDQ / QOperatoractivation_type=QuantType.QInt8, # 激活类型 Int8 / UInt8weight_type=QuantType.QInt8, # 参数类型 Int8 / UInt8calibrate_method=CalibrationMethod.MinMax, # 数据校准方法 MinMax / Entropy / Percentileper_channel=True, # 量化通道)# 可视化ONNX模型输出onnx_out_put_file = os.path.join(args.work_dir, 'output_onnxruntime_int8.png')visualize_model(model_cfg,deploy_cfg,[quant_onnx_path],args.test_img,device='cpu',backend=Backend.ONNXRUNTIME,output_file=onnx_out_put_file)#可视化torch模型输出torch_out_put_file = os.path.join(args.work_dir, 'output_torch.png')visualize_model(model_cfg,deploy_cfg,[args.checkpoint],args.test_img,device='cpu',backend=Backend.PYTORCH,output_file=torch_out_put_file)logger.info(f'torch2onnx finished. Results saved to {args.work_dir}')if __name__ == '__main__':main()
笔者任务下部分结果展示
fp32.pt模型预测结果:
onnx_int8模型的预测结果:
参考博客
MMSegmentation的用法(手把手入门教程)搭配Colab,对自己的数据进行训练
【Python】mmSegmentation语义分割框架教程(1.x版本)
欢迎指正
因为本文主要是本人用来做的笔记,顺便进行知识巩固。如果本文对你有所帮助,那么本博客的目的就已经超额完成了。
欢迎交流
邮箱:refreshmentccoffee@gmail.com