模型推理部署——Tensorflow篇

模型推理部署——Tensorflow篇

Tensorflow 模型保存、加载方式

我们常见的模型保存后产生的文件有以下几种形式：

1. checkpoint

   checkpoint \ model.ckpt-XXX.data-00000-of-00001 \ model.ckpt-XXX.index \ model.ckpt-XXX.meta

其中 .data文件是包含训练变量的文件； .index是描述variable中key和value的对应关系；checkpoint文件是列出保存的所有模型以及最近模型的相关信息；.meta文件保存了tensorflow完整的网络图结构

2. GraphDef（frozen)

   pb 文件

protobuf对象序列化后的数据,包含了计算图和变量。

3. SavedModel

   pb文件+ variable目录（.index + .data）

后面两种形式包含权重和计算图操作。它不需要运行原始的模型构建代码，这对于共享或部署（使用TFLite ， TensorFlow.js ， TensorFlow Serving或TensorFlow Hub ）很有用。

python中部署Tensorflow

和python中部署PyTorch相同，我们完全可以训练时使用python语言，推理时也使用python语言。使用flask等web开发框架，用restful API的方式部署服务。详细的可参考python中部署PyTorch中的相关内容。

Java中部署Tensorflow

因为tensorflow 提供java的API，我们可以使用python来训练模型，使用java来部署模型。

在python中保存模型pb文件（2），然后在Java中采用以下示例方式进行加载：

import java.nio.file.Files;
import java.nio.file.Path;

import org.tensorflow.Graph;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

// load pb model file 
byte[] graphDef = Files.readAllBytes(Paths.get(modelDir, "**.pb"));

Graph g = new Graph();
g.importGraphDef(graphDef);
Session model = new Session(g);

// feed input 
Tensor result = model.runner().feed("input_name",inputTensor) 

通过上述示例可以将模型加载并进行推理，可以通过http或者grpc的方式进行服务的接口通信

Tensorflow Serving

tensorflow serving 提供多模型部署、服务通信、模型管理等，我们可以使用tensorflow serving 来部署模型，撰写客户端来进行数据预处理和后处理以及与服务端的交互。

tensorflow serving 上手还是比较简单的，可以参考官方文档进行尝试使用：

https://www.tensorflow.org/tfx/guide/serving

其他

还有一些其他开源的模型部署框架服务支持tensorflow作为模型后端来进行部署。同时在推理部署时，还可以使用tensorflow-trt 或者tensorRT等进行模型操作转换，来进行推理加速。