Pytorch模型转onnx转tensorflow pb,部署tensorflow serving,client测试
Contact me
- Blog -> https://cugtyt.github.io/blog/index
- Email -> cugtyt@qq.com
- GitHub -> Cugtyt@GitHub
重点参考pytorch-onnx-tensorflow-pb和running-pytorch-models-in-production,有修改
这里假设已经有训练好的pytorch模型,完整代码在 https://cugtyt.github.io/blog/2019/torch2pb.py,需要安装pytorch,tensorflow(1.15),onnx,onnx-tf等
pytorch 转 onnx
见代码中的 export_onnx 函数
转成meta文件
见 export_tf_proto 函数
转成tensorflow serving需要的pb文件
见 export_for_serving 函数
使用docker部署tensorflow serving
官方文档见 https://www.tensorflow.org/tfx/guide/serving
文档可以启动一个docker,即运行tensorflow serving。
客户端程序
由于上一步中的文档是用一个简单的例子,因此在命令行中就可以传入数据,但是这里的输入是图像,因此编写代码进行传入。完整代码在 https://cugtyt.github.io/blog/2019/client.py。
需要说明的是这里用了grpc的api,还有tensorflow_serving(需要tensorflow 2,因此与上面的包是不兼容的)包,因此需要用8500端口,不是官方文档的8501。
docker run -p 8500:8500 --mount type=bind,source=/torch2pb/serving_model,target=/models/densenet -e MODEL_NAME=densenet -t tensorflow/serving
注意:由于参考的博客中写的是output_names = ['confidences']
,但是部署后发现tensorflow serving端报错找不到名为add的节点,其实confidences
就是指add的输出节点,这个坑不知道是什么问题。绕开的方法是output_names = ['add']
这样肯定存在add。