yolo-world 源码解析（三）（4）

# 导入必要的库
import argparse  # 用于解析命令行参数
import logging  # 用于记录日志
import os  # 用于操作系统相关功能
import os.path as osp  # 用于操作文件路径
from functools import partial  # 用于创建偏函数
import mmengine  # 导入 mmengine 库
import torch.multiprocessing as mp  # 导入多进程库
from torch.multiprocessing import Process, set_start_method  # 导入多进程相关函数
from mmdeploy.apis import (create_calib_input_data, extract_model,  # 导入一些部署相关的函数
                           get_predefined_partition_cfg, torch2onnx,
                           torch2torchscript, visualize_model)
from mmdeploy.apis.core import PIPELINE_MANAGER  # 导入核心部署管理器
from mmdeploy.apis.utils import to_backend  # 导入转换到后端的函数
from mmdeploy.backend.sdk.export_info import export2SDK  # 导入导出到 SDK 的函数
from mmdeploy.utils import (IR, Backend, get_backend, get_calib_filename,  # 导入一些工具函数
                            get_ir_config, get_partition_config,
                            get_root_logger, load_config, target_wrapper)
def parse_args():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='Export model to backends.')
    # 添加命令行参数
    parser.add_argument('deploy_cfg', help='deploy config path')  # 部署配置文件路径
    parser.add_argument('model_cfg', help='model config path')  # 模型配置文件路径
    parser.add_argument('checkpoint', help='model checkpoint path')  # 模型检查点路径
    parser.add_argument('img', help='image used to convert model model')  # 用于转换模型的图像
    parser.add_argument(
        '--test-img',
        default=None,
        type=str,
        nargs='+',
        help='image used to test model')  # 用于测试模型的图像
    parser.add_argument(
        '--work-dir',
        default=os.getcwd(),
        help='the dir to save logs and models')  # 保存日志和模型的目录
    parser.add_argument(
        '--calib-dataset-cfg',
        help='dataset config path used to calibrate in int8 mode. If not \
            specified, it will use "val" dataset in model config instead.',
        default=None)  # 用于在 int8 模式下校准的数据集配置文件路径
    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()))  # 设置日志级别
    # 添加命令行参数，用于显示检测输出
    parser.add_argument(
        '--show', action='store_true', help='Show detection outputs')
    # 添加命令行参数，用于输出 SDK 的信息
    parser.add_argument(
        '--dump-info', action='store_true', help='Output information for SDK')
    # 添加命令行参数，指定用于量化模型的图像目录
    parser.add_argument(
        '--quant-image-dir',
        default=None,
        help='Image directory for quantize model.')
    # 添加命令行参数，用于将模型量化为低位
    parser.add_argument(
        '--quant', action='store_true', help='Quantize model to low bit.')
    # 添加命令行参数，指定远程设备上进行推理的 IP 地址和端口
    parser.add_argument(
        '--uri',
        default='192.168.1.1:60000',
        help='Remote ipv4:port or ipv6:port for inference on edge device.')
    # 解析命令行参数
    args = parser.parse_args()
    # 返回解析后的参数
    return args
# 创建一个进程，执行指定的目标函数，并传入参数和关键字参数，可选地返回一个值
def create_process(name, target, args, kwargs, ret_value=None):
    # 获取根日志记录器
    logger = get_root_logger()
    # 记录进程开始信息
    logger.info(f'{name} start.')
    # 获取日志级别
    log_level = logger.level
    # 创建一个包装函数，用于处理目标函数的执行结果和日志记录
    wrap_func = partial(target_wrapper, target, log_level, ret_value)
    # 创建一个进程对象，指定目标函数、参数和关键字参数
    process = Process(target=wrap_func, args=args, kwargs=kwargs)
    # 启动进程
    process.start()
    # 等待进程结束
    process.join()
    # 如果有返回值，检查返回值是否为0，记录日志并退出程序
    if ret_value is not None:
        if ret_value.value != 0:
            logger.error(f'{name} failed.')
            exit(1)
        else:
            logger.info(f'{name} success.')
# 根据中间表示类型返回对应的转换函数
def torch2ir(ir_type: IR):
    """Return the conversion function from torch to the intermediate
    representation.
    Args:
        ir_type (IR): The type of the intermediate representation.
    """
    if ir_type == IR.ONNX:
        return torch2onnx
    elif ir_type == IR.TORCHSCRIPT:
        return torch2torchscript
    else:
        raise KeyError(f'Unexpected IR type {ir_type}')
# 主函数
def main():
    # 解析命令行参数
    args = parse_args()
    # 设置进程启动方法为'spawn'
    set_start_method('spawn', force=True)
    # 获取根日志记录器
    logger = get_root_logger()
    # 获取日志级别
    log_level = logging.getLevelName(args.log_level)
    # 设置日志级别
    logger.setLevel(log_level)
    # 定义一系列处理函数
    pipeline_funcs = [
        torch2onnx, torch2torchscript, extract_model, create_calib_input_data
    ]
    # 启用多进程模式，并指定处理函数
    PIPELINE_MANAGER.enable_multiprocess(True, pipeline_funcs)
    # 设置处理函数的日志级别
    PIPELINE_MANAGER.set_log_level(log_level, pipeline_funcs)
    # 获取部署配置路径、模型配置路径、检查点路径、量化标志和量化图片目录
    deploy_cfg_path = args.deploy_cfg
    model_cfg_path = args.model_cfg
    checkpoint_path = args.checkpoint
    quant = args.quant
    quant_image_dir = args.quant_image_dir
    # 加载部署配置和模型配置
    deploy_cfg, model_cfg = load_config(deploy_cfg_path, model_cfg_path)
    # 如果需要，创建工作目录
    mmengine.mkdir_or_exist(osp.abspath(args.work_dir))
    # 如果需要，导出到SDK
    if args.dump_info:
        export2SDK(
            deploy_cfg,
            model_cfg,
            args.work_dir,
            pth=checkpoint_path,
            device=args.device)
    # 创建一个共享变量，用于存储返回值
    ret_value = mp.Value('d', 0, lock=False)
    # 转换为中间表示
    ir_config = get_ir_config(deploy_cfg)
    # 获取保存 IR 结果的文件名
    ir_save_file = ir_config['save_file']
    # 获取 IR 类型
    ir_type = IR.get(ir_config['type'])
    # 调用 torch2ir 函数将模型转换为 IR
    torch2ir(ir_type)(
        args.img,
        args.work_dir,
        ir_save_file,
        deploy_cfg_path,
        model_cfg_path,
        checkpoint_path,
        device=args.device)
    # 转换后的 IR 文件列表
    ir_files = [osp.join(args.work_dir, ir_save_file)]
    # 获取模型分区配置
    partition_cfgs = get_partition_config(deploy_cfg)
    if partition_cfgs is not None:
        if 'partition_cfg' in partition_cfgs:
            partition_cfgs = partition_cfgs.get('partition_cfg', None)
        else:
            assert 'type' in partition_cfgs
            partition_cfgs = get_predefined_partition_cfg(
                deploy_cfg, partition_cfgs['type'])
        # 原始 IR 文件
        origin_ir_file = ir_files[0]
        ir_files = []
        # 遍历分区配置，对模型进行分区
        for partition_cfg in partition_cfgs:
            save_file = partition_cfg['save_file']
            save_path = osp.join(args.work_dir, save_file)
            start = partition_cfg['start']
            end = partition_cfg['end']
            dynamic_axes = partition_cfg.get('dynamic_axes', None)
            # 提取模型的部分并保存
            extract_model(
                origin_ir_file,
                start,
                end,
                dynamic_axes=dynamic_axes,
                save_file=save_path)
            ir_files.append(save_path)
    # 获取校准数据文件名
    calib_filename = get_calib_filename(deploy_cfg)
    if calib_filename is not None:
        calib_path = osp.join(args.work_dir, calib_filename)
        # 创建校准输入数据
        create_calib_input_data(
            calib_path,
            deploy_cfg_path,
            model_cfg_path,
            checkpoint_path,
            dataset_cfg=args.calib_dataset_cfg,
            dataset_type='val',
            device=args.device)
    # 后端文件列表
    backend_files = ir_files
    # 获取后端类型
    backend = get_backend(deploy_cfg)
    # 预处理部署配置
    # 如果选择的后端是RKNN
    if backend == Backend.RKNN:
        # TODO: 在将来将此功能添加到任务处理器中
        # 导入临时文件模块
        import tempfile
        # 从mmdeploy.utils中导入必要的函数
        from mmdeploy.utils import (get_common_config, get_normalization,
                                    get_quantization_config,
                                    get_rknn_quantization)
        
        # 获取量化配置
        quantization_cfg = get_quantization_config(deploy_cfg)
        # 获取通用配置
        common_params = get_common_config(deploy_cfg)
        
        # 如果需要RKNN量化
        if get_rknn_quantization(deploy_cfg) is True:
            # 获取归一化转换参数
            transform = get_normalization(model_cfg)
            # 更新通用参数，包括均值和标准差
            common_params.update(
                dict(
                    mean_values=[transform['mean']],
                    std_values=[transform['std']]))
        # 创建临时文件用于存储数据集文件路径
        dataset_file = tempfile.NamedTemporaryFile(suffix='.txt').name
        # 将图像文件的绝对路径写入数据集文件
        with open(dataset_file, 'w') as f:
            f.writelines([osp.abspath(args.img)])
        
        # 如果量化配置中未指定数据集，则将数据集文件路径添加到量化配置中
        if quantization_cfg.get('dataset', None) is None:
            quantization_cfg['dataset'] = dataset_file
    
    # 如果选择的后端是ASCEND
    if backend == Backend.ASCEND:
        # TODO: 在将来将此功能添加到后端管理器中
        # 如果需要输出信息
        if args.dump_info:
            # 从mmdeploy.backend.ascend中导入更新SDK管道的函数
            from mmdeploy.backend.ascend import update_sdk_pipeline
            # 更新SDK管道
            update_sdk_pipeline(args.work_dir)
    # 如果后端是VACC
    if backend == Backend.VACC:
        # TODO: 将此部分在未来添加到任务处理器中
        # 导入获取量化数据的模块
        from onnx2vacc_quant_dataset import get_quant
        # 导入获取模型输入的工具函数
        from mmdeploy.utils import get_model_inputs
        # 加载部署配置和模型配置
        deploy_cfg, model_cfg = load_config(deploy_cfg_path, model_cfg_path)
        # 获取模型输入
        model_inputs = get_model_inputs(deploy_cfg)
        # 遍历ONNX文件路径和模型输入
        for onnx_path, model_input in zip(ir_files, model_inputs):
            # 获取量化模式
            quant_mode = model_input.get('qconfig', {}).get('dtype', 'fp16')
            # 确保量化模式为'int8'或'fp16'
            assert quant_mode in ['int8', 'fp16'], quant_mode + ' not support now'
            shape_dict = model_input.get('shape', {})
            # 如果量化模式为'int8'
            if quant_mode == 'int8':
                # 创建处理过程
                create_process(
                    'vacc quant dataset',
                    target=get_quant,
                    args=(deploy_cfg, model_cfg, shape_dict, checkpoint_path,
                          args.work_dir, args.device),
                    kwargs=dict(),
                    ret_value=ret_value)
    # 设置日志级别
    PIPELINE_MANAGER.set_log_level(log_level, [to_backend])
    # 如果后端是TensorRT，启用多进程
    if backend == Backend.TENSORRT:
        PIPELINE_MANAGER.enable_multiprocess(True, [to_backend])
    # 将转换后的文件传递给后端处理
    backend_files = to_backend(
        backend,
        ir_files,
        work_dir=args.work_dir,
        deploy_cfg=deploy_cfg,
        log_level=log_level,
        device=args.device,
        uri=args.uri)
    # 进行ncnn量化
    # 如果后端为NCNN且需要量化
    if backend == Backend.NCNN and quant:
        # 导入获取量化表的函数
        from onnx2ncnn_quant_table import get_table
        # 导入NCNN相关函数
        from mmdeploy.apis.ncnn import get_quant_model_file, ncnn2int8
        # 获取模型参数路径列表
        model_param_paths = backend_files[::2]
        # 获取模型二进制文件路径列表
        model_bin_paths = backend_files[1::2]
        # 清空后端文件列表
        backend_files = []
        # 遍历ONNX文件路径、模型参数路径、模型二进制文件路径
        for onnx_path, model_param_path, model_bin_path in zip(
                ir_files, model_param_paths, model_bin_paths):
            # 加载部署配置和模型配置
            deploy_cfg, model_cfg = load_config(deploy_cfg_path,
                                                model_cfg_path)
            # 获取量化后的ONNX文件、量化表、量化参数、量化二进制文件
            quant_onnx, quant_table, quant_param, quant_bin = get_quant_model_file(
                onnx_path, args.work_dir)
            # 创建进程，获取量化表
            create_process(
                'ncnn quant table',
                target=get_table,
                args=(onnx_path, deploy_cfg, model_cfg, quant_onnx,
                      quant_table, quant_image_dir, args.device),
                kwargs=dict(),
                ret_value=ret_value)
            # 创建进程，进行NCNN量化
            create_process(
                'ncnn_int8',
                target=ncnn2int8,
                args=(model_param_path, model_bin_path, quant_table,
                      quant_param, quant_bin),
                kwargs=dict(),
                ret_value=ret_value)
            # 将量化参数和量化二进制文件添加到后端文件列表中
            backend_files += [quant_param, quant_bin]
    # 如果未指定测试图片，则使用默认图片
    if args.test_img is None:
        args.test_img = args.img
    # 额外参数，包括后端类型、输出文件路径、是否显示结果
    extra = dict(
        backend=backend,
        output_file=osp.join(args.work_dir, f'output_{backend.value}.jpg'),
        show_result=args.show)
    # 如果后端为SNPE，则添加URI参数
    if backend == Backend.SNPE:
        extra['uri'] = args.uri
    # 获取后端推理结果，并尝试渲染
    create_process(
        f'visualize {backend.value} model',
        target=visualize_model,
        args=(model_cfg_path, deploy_cfg_path, backend_files, args.test_img,
              args.device),
        kwargs=extra,
        ret_value=ret_value)
    # 获取PyTorch模型推理结果，尝试可视化（如果可能）
    # 创建一个进程，用于可视化 PyTorch 模型
    create_process(
        'visualize pytorch model',  # 进程名称
        target=visualize_model,  # 目标函数为 visualize_model
        args=(model_cfg_path, deploy_cfg_path, [checkpoint_path], args.test_img, args.device),  # 参数列表
        kwargs=dict(
            backend=Backend.PYTORCH,  # 使用 PyTorch 后端
            output_file=osp.join(args.work_dir, 'output_pytorch.jpg'),  # 输出文件路径
            show_result=args.show),  # 是否显示结果
        ret_value=ret_value)  # 返回值为 ret_value
    # 记录信息到日志
    logger.info('All process success.')
# 如果当前脚本被直接执行，则调用主函数
if __name__ == '__main__':
    main()

# 导入必要的库
import argparse
import os.path as osp
from copy import deepcopy
# 导入自定义模块
from mmengine import DictAction
from mmdeploy.apis import build_task_processor
from mmdeploy.utils.config_utils import load_config
from mmdeploy.utils.timer import TimeCounter
# 解析命令行参数
def parse_args():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='MMDeploy test (and eval) a backend.')
    
    # 添加命令行参数
    parser.add_argument('deploy_cfg', help='Deploy config path')
    parser.add_argument('model_cfg', help='Model config path')
    parser.add_argument('--model', type=str, nargs='+', help='Input model files.')
    parser.add_argument('--device', help='device used for conversion', default='cpu')
    parser.add_argument('--work-dir', default='./work_dir', help='the directory to save the file containing evaluation metrics')
    parser.add_argument('--cfg-options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair in xxx=yyy format will be merged into config file. If the value to be overwritten is a list, it should be like key="[a,b]" or key=a,b It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" Note that the quotation marks are necessary and that no white space is allowed.')
    parser.add_argument('--show', action='store_true', help='show results')
    parser.add_argument('--show-dir', help='directory where painted images will be saved')
    parser.add_argument('--interval', type=int, default=1, help='visualize per interval samples.')
    parser.add_argument('--wait-time', type=float, default=2, help='display time of every window. (second)')
    parser.add_argument('--log2file', type=str, help='log evaluation results and speed to file', default=None)
    # 添加命令行参数，用于激活速度测试
    parser.add_argument(
        '--speed-test', action='store_true', help='activate speed test')
    # 添加命令行参数，用于设置预热次数，在计算推理时间之前需要设置速度测试
    parser.add_argument(
        '--warmup',
        type=int,
        help='warmup before counting inference elapse, require setting '
        'speed-test first',
        default=10)
    # 添加命令行参数，用于设置日志输出间隔，在设置速度测试之前需要设置
    parser.add_argument(
        '--log-interval',
        type=int,
        help='the interval between each log, require setting '
        'speed-test first',
        default=100)
    # 添加命令行参数，用于设置测试的批次大小，会覆盖数据配置中的 `samples_per_gpu`
    parser.add_argument(
        '--batch-size',
        type=int,
        default=1,
        help='the batch size for test, would override `samples_per_gpu`'
        'in  data config.')
    # 添加命令行参数，用于设置远程推理设备的 IP 地址和端口
    parser.add_argument(
        '--uri',
        action='store_true',
        default='192.168.1.1:60000',
        help='Remote ipv4:port or ipv6:port for inference on edge device.')
    # 解析命令行参数
    args = parser.parse_args()
    # 返回解析后的参数
    return args
def main():
    # 解析命令行参数
    args = parse_args()
    # 获取部署配置文件路径和模型配置文件路径
    deploy_cfg_path = args.deploy_cfg
    model_cfg_path = args.model_cfg
    # 加载部署配置和模型配置
    deploy_cfg, model_cfg = load_config(deploy_cfg_path, model_cfg_path)
    # 确定工作目录的优先级：命令行参数 > 文件中的段落 > 文件名
    if args.work_dir is not None:
        # 如果命令行参数中指定了工作目录，则更新配置
        work_dir = args.work_dir
    elif model_cfg.get('work_dir', None) is None:
        # 如果配置文件中未指定工作目录，则使用配置文件名作为默认工作目录
        work_dir = osp.join('./work_dirs', osp.splitext(osp.basename(args.config))[0])
    # 合并模型配置的选项
    if args.cfg_options is not None:
        model_cfg.merge_from_dict(args.cfg_options)
    # 构建任务处理器
    task_processor = build_task_processor(model_cfg, deploy_cfg, args.device)
    # 准备数据集加载器
    test_dataloader = deepcopy(model_cfg['test_dataloader'])
    if isinstance(test_dataloader, list):
        dataset = []
        for loader in test_dataloader:
            # 构建数据集
            ds = task_processor.build_dataset(loader['dataset'])
            dataset.append(ds)
            loader['dataset'] = ds
            loader['batch_size'] = args.batch_size
            loader = task_processor.build_dataloader(loader)
        dataloader = test_dataloader
    else:
        test_dataloader['batch_size'] = args.batch_size
        dataset = task_processor.build_dataset(test_dataloader['dataset'])
        test_dataloader['dataset'] = dataset
        dataloader = task_processor.build_dataloader(test_dataloader)
    # 加载后端模型
    model = task_processor.build_backend_model(
        args.model,
        data_preprocessor_updater=task_processor.update_data_preprocessor)
    destroy_model = model.destroy
    is_device_cpu = (args.device == 'cpu')
    # 使用任务处理器构建测试运行器，传入模型、工作目录、是否记录日志到文件、是否展示测试结果、展示目录、等待时间、间隔时间、数据加载器等参数
    runner = task_processor.build_test_runner(
        model,
        work_dir,
        log_file=args.log2file,
        show=args.show,
        show_dir=args.show_dir,
        wait_time=args.wait_time,
        interval=args.interval,
        dataloader=dataloader)
    # 如果需要进行速度测试
    if args.speed_test:
        # 根据设备是否为 CPU 决定是否需要同步
        with_sync = not is_device_cpu
        # 激活时间计数器，设置预热次数、日志间隔、是否同步、记录日志到文件、批处理大小等参数
        with TimeCounter.activate(
                warmup=args.warmup,
                log_interval=args.log_interval,
                with_sync=with_sync,
                file=args.log2file,
                batch_size=args.batch_size):
            # 运行测试
            runner.test()
    else:
        # 运行测试
        runner.test()
    
    # 仅在后端需要显式清理时生效（例如 Ascend）
    destroy_model()
# 如果当前脚本被直接执行，则调用主函数
if __name__ == '__main__':
    main()