MMsegmentation与MMdeploy简单使用

最近涉及到了图像分割的任务，于是拿来写下博客加深下使用。
MMsegmentation与MMdeploy的环境配置暂不做讲解，在官网和其他博客中有很多说明。
MMdeploy主要是把pt转为 onnx_int8的情况。
MMsegmentation环境配置可以参考 : 安装与配置MMSegmentation

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

$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_processor


def 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 None
        
        img_path = data['data_samples'].img_path
        
        data, 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 args


def 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_cfg
    deploy_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 onnx
        from onnxconverter_common import float16

        common_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 / QOperator
                    activation_type=QuantType.QInt8, # 激活类型 Int8 / UInt8
                    weight_type=QuantType.QInt8, # 参数类型 Int8 / UInt8
                    calibrate_method=CalibrationMethod.MinMax, # 数据校准方法 MinMax / Entropy / Percentile
                    per_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