这篇教程介绍了如何编写模型转换和部署的配置文件。部署配置文件由ONNX配置
, 代码库配置
, 推理框架配置
组成。
ONNX 配置描述了如何将PyTorch模型转换为ONNX模型。
-
type
: 配置类型。 默认为onnx
。 -
export_params
: 如果指定,将导出模型所有参数。如果您只想导出未训练模型将此项设置为 False。 -
keep_initializers_as_inputs
: 如果为 True,则所有初始化器(通常对应为参数)也将作为输入导出,添加到计算图中。 如果为 False,则初始化器不会作为输入导出,不添加到计算图中,仅将非参数输入添加到计算图中。 -
opset_version
: ONNX的算子集版本,默认为11。 -
save_file
: 输出ONNX模型文件。 -
input_names
: 模型计算图中输入节点的名称。 -
output_names
: 模型计算图中输出节点的名称。 -
input_shape
: 模型输入张量的高度和宽度。
onnx_config = dict(
type='onnx',
export_params=True,
keep_initializers_as_inputs=False,
opset_version=11,
save_file='end2end.onnx',
input_names=['input'],
output_names=['output'],
input_shape=None)
如果模型要求动态尺寸的输入和输出,您需要在ONNX配置中加入dynamic_axes配置。
dynamic_axes
: 描述输入和输出的维度信息。
dynamic_axes={
'input': {
0: 'batch',
2: 'height',
3: 'width'
},
'dets': {
0: 'batch',
1: 'num_dets',
},
'labels': {
0: 'batch',
1: 'num_dets',
},
}
代码库主要指OpenMMLab 系列模型代码库,代码库配置由OpenMMLab 系列模型代码库的简称和OpenMMLab 系列模型任务类型组成。
type
: OpenMMLab 系列模型代码库的简称, 包括mmpretrain
,mmdet
,mmseg
,mmocr
,mmagic
。task
: OpenMMLab 系列模型任务类型, 具体请参考 OpenMMLab 系列模型任务列表。
codebase_config = dict(type='mmpretrain', task='Classification')
推理框架配置主要用于指定模型运行在哪个推理框架,并提供模型在推理框架运行时所需的信息,具体参考 ONNX Runtime, TensorRT, ncnn, PPLNN。
type
: 模型推理框架, 包括onnxruntime
,ncnn
,pplnn
,tensorrt
,openvino
。
backend_config = dict(
type='tensorrt',
common_config=dict(
fp16_mode=False, max_workspace_size=1 << 30),
model_inputs=[
dict(
input_shapes=dict(
input=dict(
min_shape=[1, 3, 512, 1024],
opt_shape=[1, 3, 1024, 2048],
max_shape=[1, 3, 2048, 2048])))
])
这里我们提供了一个以TensorRT为推理框架的基于mmpretrain图像分类任务的完整部署配置示例。
codebase_config = dict(type='mmpretrain', task='Classification')
backend_config = dict(
type='tensorrt',
common_config=dict(
fp16_mode=False,
max_workspace_size=1 << 30),
model_inputs=[
dict(
input_shapes=dict(
input=dict(
min_shape=[1, 3, 224, 224],
opt_shape=[4, 3, 224, 224],
max_shape=[64, 3, 224, 224])))])
onnx_config = dict(
type='onnx',
dynamic_axes={
'input': {
0: 'batch',
2: 'height',
3: 'width'
},
'output': {
0: 'batch'
}
},
export_params=True,
keep_initializers_as_inputs=False,
opset_version=11,
save_file='end2end.onnx',
input_names=['input'],
output_names=['output'],
input_shape=[224, 224])
我们遵循以下样式来命名配置文件。建议贡献者遵循相同的风格。
(task name)_(backend name)_(dynamic or static).py
task name
: 模型任务类型。backend name
: 推理框架名称。注意:如果您使用了量化,您需要指出量化类型。例如tensorrt-int8
。dynamic or static
: 动态或者静态尺寸导出。 注意:如果推理框架需要明确的形状信息,您需要添加输入大小的描述,格式为高度 x 宽度
。 例如dynamic-512x1024-2048x2048
, 这意味着最小输入形状是512x1024
,最大输入形状是2048x2048
。
detection_tensorrt-int8_dynamic-320x320-1344x1344.py
请根据模型具体任务的代码库,编写模型配置文件。 模型配置文件用于初始化模型,详情请参考MMPretrain,MMDetection, MMSegmentation, MMOCR,MMagic。