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 prepare
model = 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"
}
}
}
}
}