ANOMALIB 第三章：模型转换

ANOMALIB 第三章：模型转换

前两章我们部署了anomalib项目并尝试训练了一个模型，最后得到一个lighting pytorch格式的ckpt权重文件，并且anomalib平台支持将模型转为openvino等格式的文件。从训练到部署非常方便，但这是基于openvino平台而言（毕竟本身就是openvino的第一方项目）。然而基于intel cpu的openvino实在太慢了（相信大部分人都没有arch gpu），而nvidia的gpu才是模型推理的最终答案。怎么样把ckpt转为最终的engine呢？
其实anomalib是支持直接将ckpt转为onnx的，再将onnx转为engine不就好了。（总感觉背叛了anomalib，很抱歉和（；´д｀）ゞ）

第一步： ckpt转onnx

仔细阅读anomalib的官方文档（非常粗糙的文档），我们可以发现官网是给出了转换的功能的。然后自行阅读代码，我们可以找到对应的转换接口，尝试写个代码调用一下：

from anomalib.models import Patchcore
from anomalib.engine import Engine

model = Patchcore()
engine = Engine(task="classification")
onnx_model = engine.export(
    model=model,
    export_type='onnx',
    export_root=None,
    input_size=[244, 244],
    transform=None,
    compression_type=None,
    datamodule=None,
    metric=None,
    ov_args=None,
    ckpt_path='E:\\proj\\anomalib\\myProj\\model.ckpt',  # 存放model.ckpt文件的路径，需要对应修改
)
print(onnx_model)

这里解释一下代码流程，首先我们要创建一个模型对象并定义模型任务，然后我们创建一个模型的导出对象。要输入的形参分别是：

1. model:模型、
2. export_type:导出类型、
3. export_root: 导出根目录、
4. input_size: 定义推理图像的输入宽高、
5. transform: 预处理、
6. compression_type: 用于指定导出模型时是否进行压缩以及压缩类型、
7. datamodule: 如果导出过程中需要特定的数据集或数据处理逻辑，可以通过这个参数传递一个数据模块用于评估模型精度
8. metric： 用于指定评估模型性能的度量标准
9. ov_args： 用于传递额外的参数给OpenVINO优化工具
   接着我们准备好模型，直接运行应该就能在results的weights文件夹下看到onnx模型。

第二步： 将onnx模型转为tensorrt的engine模型

这一步就很简单了，网上也一堆教程和原理解释，我就直接贴代码，实测patchcore的onnx转engine可用，如果遇到问题欢迎讨论。

import numpy as np
import tensorrt as trt
import torch
import logging
from tensorrt import IInt8Calibrator

# logger to capture errors, warnings, and other information during the build and inference phases
TRT_LOGGER = trt.Logger()


def build_engine(onnx, dynamic=True, half=True):
   # f = onnx.with_suffix('.engine')
   f = 'patchcore_shape_test.engine'
   # 1、创建日志记录器
   log = trt.Logger()
   # 2、创建builder对象
   builder = trt.Builder(log)
   # 3、创建 Builder Config 对象
   config = builder.create_builder_config()
   # 4、将workspace*1 二进制左移30位后的10进制
   workspace = 1
   config.max_workspace_size = workspace * 1 << 30  # 设置 TensorRT 推理引擎使用的最大工作空间大小，单位为字节
   # 5、定义networko并加载ONNX解析器
   flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
   network = builder.create_network(flag)
   parser = trt.OnnxParser(network, log)

   if not parser.parse_from_file(str(onnx)):  # 查看是否解析成功
       raise RuntimeError(f'failed to load ONNX file: {onnx}')

   # 6、获得网络的输入输出
   inputs = [network.get_input(i) for i in range(network.num_inputs)]
   outputs = [network.get_output(i) for i in range(network.num_outputs)]

   # 7.判断是否动态输入
   if dynamic:
       im = torch.zeros(1, 3, 224, 224)
       if im.shape[0] <= 1:
           # log.warning(f"{trt} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
           print('x')
       profile = builder.create_optimization_profile()
       for inp in inputs:
           profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
       config.add_optimization_profile(profile)
   # 判断是否支持FP16推理

   if builder.platform_has_fast_fp16 and half:
       config.set_flag(trt.BuilderFlag.FP16)
   # build engine 文件的写入  这里的f是前面定义的engine文件
   with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
       # 序列化model
       t.write(engine.serialize())
   return f, None


if __name__ == '__main__':
   engine, context = build_engine(r'D:\\proj\\resource\\model\\patchcore.onnx')

这样我们就能得到一个推理速度超快的模型啦。主要这本来是一件非常普通且正常的事情，奈何官网文档实在过于简陋了，所以才想着特别记录一下。