1.onnx runtime安装
# 激活虚拟环境 conda activate env_name # env_name换成环境名称 # 安装onnx pip install onnx # 安装onnx runtime pip install onnxruntime # 使用CPU进行推理 # pip install onnxruntime-gpu # 使用GPU进行推理
2.导出模型
import torch.onnx # 转换的onnx格式的名称，文件后缀需为.onnx onnx_file_name = "xxxxxx.onnx" # 我们需要转换的模型，将torch_model设置为自己的模型 model = torch_model # 加载权重，将model.pth转换为自己的模型权重 # 如果模型的权重是使用多卡训练出来，我们需要去除权重中多的module. 具体操作可以见5.4节 model = model.load_state_dict(torch.load("model.pth")) # 导出模型前，必须调用model.eval()或者model.train(False) model.eval() # dummy_input就是一个输入的实例，仅提供输入shape、type等信息 batch_size = 1 # 随机的取值，当设置dynamic_axes后影响不大 dummy_input = torch.randn(batch_size, 1, 224, 224, requires_grad=True) # 这组输入对应的模型输出 output = model(dummy_input) # 导出模型 torch.onnx.export(model, # 模型的名称 dummy_input, # 一组实例化输入 onnx_file_name, # 文件保存路径/名称 export_params=True, # 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False. opset_version=10, # ONNX 算子集的版本，当前已更新到15 do_constant_folding=True, # 是否执行常量折叠优化 input_names = ['input'], # 输入模型的张量的名称 output_names = ['output'], # 输出模型的张量的名称 # dynamic_axes将batch_size的维度指定为动态， # 后续进行推理的数据可以与导出的dummy_input的batch_size不同 dynamic_axes={'input' : {0 : 'batch_size'}, 'output' : {0 : 'batch_size'}})
3.模型校验
import onnx # 我们可以使用异常处理的方法进行检验 try: # 当我们的模型不可用时，将会报出异常 onnx.checker.check_model(self.onnx_model) except onnx.checker.ValidationError as e: print("The model is invalid: %s"%e) else: # 模型可用时，将不会报出异常，并会输出“The model is valid!” print("The model is valid!")
4.模型可视化
Netron下载网址：github.com/lutzroeder/…
5.使用ONNX Runtime进行推理
使用ONNX Runtime运行一下转化后的模型，看一下推理后的结果。
# 导入onnxruntime import onnxruntime # 需要进行推理的onnx模型文件名称 onnx_file_name = "xxxxxx.onnx" # onnxruntime.InferenceSession用于获取一个 ONNX Runtime 推理器 ort_session = onnxruntime.InferenceSession(onnx_file_name) # 构建字典的输入数据，字典的key需要与我们构建onnx模型时的input_names相同 # 输入的input_img 也需要改变为ndarray格式 ort_inputs = {'input': input_img} # 我们更建议使用下面这种方法,因为避免了手动输入key # ort_inputs = {ort_session.get_inputs()[0].name:input_img} # run是进行模型的推理，第一个参数为输出张量名的列表，一般情况可以设置为None # 第二个参数为构建的输入值的字典 # 由于返回的结果被列表嵌套，因此我们需要进行[0]的索引 ort_output = ort_session.run(None,ort_inputs)[0] # output = {ort_session.get_outputs()[0].name} # ort_output = ort_session.run([output], ort_inputs)[0]
