显微镜图像处理【部署】- pytorch模型转onnx使用GPU进行推理

文章不涉及具体的算法训练，但是如果你不了解如何用pytorch训练一个最简单的图像预测模型的话，我可以推荐你阅读这篇文章：github.com/open-mmlab/...
在处理显微镜图像时，我们会有各式各样的深度学习或机器学习算法部署需求。你可能需要对显微镜拍摄出来的图像进行对你的业务而言有价值的图像数据分析。如果你是显微镜图像处理特定领域的专业人士的话。
下面以核型分析业务为例，训练一个pytorch的pth模型文件，并将该模型记录静态图，转为onnx模型进行推理调用。并分析onnx模型对于软件项目在windows10操作系统上部署的性能提升。

性能评价指标

对于一个windows系统而言，性能的评价指标应该分为以下多个方面：

软件构建产物体积：对于附带深度学习或机器学习模块的软件项目而言，部署的体积问题是开发者应该重点关注的部分，应用体积的提升会导致软件推广产生困难，一个附带pyqt或者pytorch功能模块的python软件项目打包之后产生的体积是恐怖的，就算你使用nuitka进行打包，将所有的依赖包构建为C语言模块，也会生成一个至少1.5GB的最终软件产物。如果你将所有的依赖包都复制进最终打包的产物，而不是构建，那么软件体积将达到可怕的5~10GB【单torch模块就要占用3GB】。这显然是不可以接受的
算法运行效率：一个附带深度学习或者机器学习模块的软件项目而言，算法构建产物的管理是开发者需要去考虑的非常痛苦的一个工作模块，以5000张图用循环神经网络跑10个epoch25个class的实例分割模型为例，最终的pth构建产物会在500MB左右，并且你的项目可能会存在多个深度学习模块，比如图像风格化迁移、图像超分辨率、图像分类、图像打分、图像数据评价模块。这些模块的加载调用和预热非常占用GPU内存。如果你的机器内存只有16GB或者8GB，或者你的显卡显存只有8GB或4GB,同时运行这么多深度学习模块是不现实的，这样就增加了机器的运行成本以及阻碍了软件及时的向你的客户进行推广【我的软件销售风格是：软件下载了就走，走了就用，用了没bug，减少客户心智负担】
软件架构稳定性：你的作品应该在用户使用的过程中尽可能的没有瑕疵，这个是必须的。要减少用户心智负担，那么就代表着软件的启动速度，算法模块调用合理性，UI界面的交互和操作美观性，都是必须要应用上best practiecs的。

而使用onnx模型替代原始pth文件，是必要的软件架构升级项目，是新时代AI杀手应用的入场券。

onnx打包产物结构分析

首先推荐一个网站：netron.app/ 这个网站能够上传你的onnx模型，并且分析你的onnx模型的graph。你不管是自己烘培的pytorch模型或者是你从网上download下来的其他人的pytorch模型，你都能够通过这个网站查看你的onnx模型的输入和输出【这个可能是starter最关注的部分】

图1【onnx结构分析】

加载onnx模型

python 复制代码

import onnx
model = onnx.load("work_dir/onnx/cascade_rcnn/end2end.onnx")
graph = model.graph
node = graph.node
input = graph.input
output = graph.output

node: onnx模型是一个pth模型推理一遍过后记录下来的静态图，node代表神经网络中推理会运行的每一个步骤
input: onnx模型的输入参数
output: onnx模型的输出参数

根据图1可知，本篇文章调用的onnx模型接受一个叫做input的输入，接受一张tensor图像，图像支持任意尺寸，输出三个参数，分别为det，masks,labels。

调用onnx模型

python 复制代码

def preprocess_input(img_path):
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 保持原始分辨率
    img = img.astype(np.float32) / 255.0        # 自动适应任意尺寸
    return img.transpose(2, 0, 1)[None,...]     # HWC -> NCHW
    
input_img = preprocess_input1("img_4.png")

ort_session = onnxruntime.InferenceSession("work_dir/onnx/cascade_rcnn/end2end.onnx")

ort_inputs = {'input': input_img}

labels,masks,dets = ort_session.run(['labels','masks','dets'], ort_inputs)

通过以上代码能够调用onnx对一张图片进行推理，并输出推理结果

onnx推理结果可视化

图2【传入推理的图片】

图3【推理结果可视化】