Pytorch模型的TF-Serving部署
Pytorch模型的TF-Serving部署
写在前⾯
⽬前PyTorch在学术界⼏乎完全盖过了Tensorflow，从个⼈体验来说：
1. PyTorch并没有表现出⽐Tensorflow性能差很多，基本是相当；
2. 与Tensorflow的⾮Eager API相⽐，PyTorch上⼿、使⽤的难度相对来说⾮常低，并且调试⾮常容易（可断点可单步）；
3. 以及笔者还没有体会到的动态图与静态图之争
但是在应⽤场景下仍然还有将模型使⽤TF-Serving部署的需求，做转换还是有意义的。
本⽂将以⼀个transformer为例来介绍整个流程。
1. PyTorch侧，利⽤onnx导出
依赖：onnx
model = MultiTaskModel(encoder, sst_head, mrpc_head)
# 关注点1
device = torch.device('cuda'if torch.cuda.is_available()else'cpu')
# 关注点2
tensor = sor((), dtype=torch.int64)
input_ids = w_ones((1,48)).to(device)
token_type_ids = w_ones((1,48)).to(device)
attention_mask = w_ones((1,48)).to(device)
# 关注点3
(input_ids, attention_mask, token_type_ids),
"/path/",
input_names=["input_ids","attention_mask","token_type_ids"],
output_names=["intent","predicate"],
dynamic_axes={"input_ids":[0],"attention_mask":[0],"token_type_ids":[0],"intent":[0],"predicate":[0]})
关注点：
1. onnx导出的时候需要给模型提供虚拟数据（input_ids/token_type_ids/attention_mask均是），数值不重要，也不⽤担⼼batch设为1了，
后续模型就真的只能接收⼤⼩是(1, 48)的参数了（固定参数不过仍然有调整空间）
2. .to(device)如果都是在CPU上的时候可以不设，但是在GPU上就需要显式地让模型和参数位于相同设备
3. port参数解释：
1. 第⼆个参数是元组，与model在forward接收的参数⼀致，所以将model的参数平铺（避开Map/List等
等）会让导出更直接，
也可以把原模型中的组件抽出来新构建⼀个模型（⽐如多任务学习的框架⼀般很复杂，但是可以⼿⼯构造⼀个简单的，把
encoder和⽤到的head⼀起作为简单参数）；
2. dynamic_axes：设置之后相应的维度就会是可变的，有需要可以设（这⾥是将batch_size设置为了可变）在明确单点预测的情
况下可以忽略。
3. opset_version: 使⽤默认值就好了，使⽤更⾼的值在下游转换时可能出现opset未实现的错误
1. 2021/03/13 补充：如果导出⽬标是tensorflow
2.x，opset_version可以设为10
2. Tensorflow侧额外两步
依赖 onnx-tf
onnx转换出来的单pb模型不好部署（⽐如Java本地加载调试，笔者只⾛通了saved_model），需要再转换为saved_model。
其实看看代码，直接转换⼀步到位应该也是可以的，没有去探索偷个懒。
2.1 onnx转单pb模型
可以使⽤命令
onnx-tf convert -- --outfile model.pb
也可以上代码，⽅便集成
import onnx
from onnx_tf.backend import prepare
onnx_model = onnx.load("/")# load onnx model
tf_rep = prepare(onnx_model)
port_graph("/path/model.pb")
2.2 单pb转saved_model
2.1 中转出的单pb模型没有签名，加载有些⿇烦，重新加载加签名加标签。
import tensorflow as tf
import os
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
with tf.gfile.GFile("/path/model.pb","rb")as f:
graph_def = tf.GraphDef()
graph_def.ad())
export_dir ="/path/saved_model"
if ists(export_dir):# 不会⾃动覆盖，需要⼿⼯删除
sigs ={}
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
with tf.Session(graph=tf.Graph())as sess:
# name="" is important to ensure we don't get spurious prefixing
tf.import_graph_def(graph_def, name="")
tf_graph = tf.get_default_graph()
input_ids = _tensor_by_name("input_ids:0")
attention_mask = _tensor_by_name("attention_mask:0")
token_type_ids = _tensor_by_name("token_type_ids:0")
intent = _tensor_by_name("intent:0")
predicate = _tensor_by_name("predicate:0")
sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]= \
tf.saved_model.signature_def_utils.predict_signature_def(
{"input_ids": input_ids,"attention_mask": attention_mask,"token_type_ids": token_type_ids},tensorflow版本选择
{"intent": intent,"predicate": predicate})
builder.add_meta_graph_and_variables(sess,
[tag_constants.SERVING],
signature_def_map=sigs)
这样看，TF2.0不像1.x有各种显⽰的变量名称，在保存模型时也这样来⼀次或者转换⼀次、打上签名和标签，能够让调⽤更简单。最终会⽣成⽂件和⽬录如下：
saved_model
|-- variables
|-- saved_model.pb
2.3 可能遇到的问题：
INVALID_ARGUMENT: NodeDef mentions attr ‘incompatible_shape_error’ not in Op<name=NotEqual; signat ure=x:T, y:T -> z:bool; attr=T:type,allowed=[DT_BFLOAT16, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_UINT8, …, DT_QINT8,
DT_QINT32, DT_STR ING, DT_BOOL, DT_COMPLEX128]; is_commutative=true>; NodeDef: {{node NotEqual}}. (Check whether your GraphDef-interpreting binary i s up to date with your GraphDef-generating binary.)
这个问题折腾了我有点时间，如图加粗，还是应该好好看错误信息才对。当时使⽤的是1.15做的转换、导出，线上部署是1.14的。换到1.14上导出就能够避开这个问题了。
划重点：onnx转换时⽤的tensorflow版本最好是与线上部署使⽤的版本保持⼀致
2.4 针对 TF 2.x
本⼈在测试FP16导出时发现了TF1.14 很难到⼀个合适的onnx/onnx-tf配置，会出现各种错误，需要同步做以下更改。上⽂中并未给出配置信息（抱歉），这⾥给出⼀套验证过的配置：
pytorch==1.7.1
tensorflow==2.1
tensorflow-addon==0.9.1 # 需要与tensorflow版本同步变化
onnx==1.8.0
onnx-tf==1.7.0
1. opset_version需要设为10：有⽐较明显的错误提⽰要这么做
2. 不再需要先转成单pb⽂件再转换到saved_model，⾼版本的onnx-tf可以⼀步到位（可以通过查看导出的输⼊变量绑定，变长维度确
认），所以代码可以简化为
from sys import argv
import onnx
from onnx_tf.backend import prepare
if __name__ =='__main__':
path = argv[1]
print("onnx => tensorflow saved model")
graph_pb = f"{path}/saved_model"
onnx_model = onnx.load(f"{path}/")# load onnx model
tf_rep = prepare(onnx_model)# run the loaded model
port_graph(graph_pb)
备注：已知这样操作后输出的名字丢失了（输⼊没问题），⽬前还不知道怎么前回来。
3. ⼩结
本⽂整合了⽹上的⼀些材料，对PyTorch模型转换到TF-Serving做了⼀点探索，同时也修正了⼀些⽹上的排名靠前的⽂章中的错误，还列举了⼀些遇到的其它问题。
想想验证完了所有流程，最后遇到2.3时的各种怀疑⼈⽣……嗯，做到这⾥，才算是能够愉快地使⽤PyTorch了，真⾹。
参考

版权声明：本站内容均来自互联网，仅供演示用，请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系QQ:729038198，我们将在24小时内删除。