yolov5-pytorch模型转onnx模型再转tensorflow savedmode模型进行tfserving部署

pytorch框架下的yolov5模型部署到tensorflow serving，需要将pytorch的pt模型转换为onnx模型，再将onnx模型转换为tfserving的savedmode模型。
1. pytorch的pt模型转onnx模型
使用yolov5中自带的 export.py 脚本（在models下）可以将模型导出为 TorchScript， ONNX, CoreML。
环境： yolov5-5.0的requirements.txt dependencies, including Python>=3.8 and PyTorch==1.7。

转换模型需要的包：
pip install coremltools=4.1 （>=4.1）
pip install onnx=1.9.0 （>=1.9.0）
pip install scikit-learn==0.19.2
pip install onnx-tf (最新版本，我的是1.8.0)
pip install tensorflow-gpu (我安装的最新版本的，cpu版本应该也可以)
注意：以上依赖包版本不对可能会报错，我的之前在onnx转pb模型时一直报错：NotImplementedError: Constant version 12 is not implemented.

yolov5 pt转换onnx的命令如下：

python models/export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1

此命令将预训练的YOLOv5s模型导出为ONNX，TorchScript和CoreML格式。 yolov5s.pt是最小，最快的模型。 其他选项是yolov5m.pt，yolov5l.pt和yolov5x.pt，或者您通过训练自定义数据集run / exp0 / weights / best.pt拥有检查点。命令如下：

 python models/export.py --weights runs/train/exp/weights/best.pt --img 640 --batch 1  # export at 640x640 with batch size 1

输出过程如下：

转换结果如下：

2. 使用Netron来查看你的ONNX文件模型
netron现已支持大部分格式的模型文件，都是可以查看的。
PaddlePaddle、OpenVINO、TensorFlow 、Caffe…
安装netron

pip install netron

进入python，运行netron.start(‘模型路径’)，如下：

import netron
netron.start('best.onnx')

然后复制红框中的链接浏览器打开即可看到模型的网络结构，如下：

3. onnx模型转换为tfserving的savedmode模型
在模型路径下创建一个onnx2pb.py脚本，转换代码如下：

import onnx
import numpy as np
from onnx_tf.backend import preparemodel = onnx.load('best.onnx')      # yolov5 pt模型转换得到的onnx模型
tf_model = prepare(model)
tf_model.export_graph('yolov5_saved_model')     # onnx模型转换为tfserving的savedmode模型

运行python onnx2pb.py即可，转换结果如下：


导出模型可直接用于 tensorflow_server, 签名默认 default_serving, 输入输出如下:

"inputs": [{
    'node_name': 'images', 'node_type': 'DT_FLOAT', 'node_shape': [1, 3, 640, 640]}],
"outputs": [{
    'node_name': 'output_0', 'node_type': 'DT_FLOAT', 'node_shape': [1, 3, 20, 20, 8]}]

下面是savemodel pb模型的输入、输出节点的具体信息：

"metadata": {
    "signature_def": {
    "signature_def": {
    "__saved_model_init_op": {
    "inputs": {
    },"outputs": {
    "__saved_model_init_op": {
    "dtype": "DT_INVALID","tensor_shape": {
    "dim": [],"unknown_rank": true},"name": "NoOp"}},"method_name": ""},"serving_default": {
    "inputs": {
    "images": {
    "dtype": "DT_FLOAT","tensor_shape": {
    "dim": [{
    "size": "1","name": ""},{
    "size": "3","name": ""},{
    "size": "640","name": ""},{
    "size": "640","name": ""}],"unknown_rank": false},"name": "serving_default_images:0"}},"outputs": {
    "output_0": {
    "dtype": "DT_FLOAT","tensor_shape": {
    "dim": [{
    "size": "1","name": ""},{
    "size": "3","name": ""},{
    "size": "80","name": ""},{
    "size": "80","name": ""},{
    "size": "8","name": ""}],"unknown_rank": false},"name": "StatefulPartitionedCall:0"},"output_1": {
    "dtype": "DT_FLOAT","tensor_shape": {
    "dim": [{
    "size": "1","name": ""},{
    "size": "3","name": ""},{
    "size": "40","name": ""},{
    "size": "40","name": ""},{
    "size": "8","name": ""}],"unknown_rank": false},"name": "StatefulPartitionedCall:1"},"output_2": {
    "dtype": "DT_FLOAT","tensor_shape": {
    "dim": [{
    "size": "1","name": ""},{
    "size": "3","name": ""},{
    "size": "20","name": ""},{
    "size": "20","name": ""},{
    "size": "8","name": ""}],"unknown_rank": false},"name": "StatefulPartitionedCall:2"}},"method_name": "tensorflow/serving/predict"}}
}
}
}