Pytorch模型的TF-Serving部署

最新推荐文章于 2024-08-09 07:32:48 发布

du00

最新推荐文章于 2024-08-09 07:32:48 发布

阅读量5.5k

点赞数 9

分类专栏：工具文章标签： tensorflow pytorch

本文链接：https://blog.youkuaiyun.com/duh2so4/article/details/109849549

版权

本文介绍了如何将PyTorch模型通过ONNX转换为TensorFlow的saved_model，以适应TF-Serving部署。过程中涉及PyTorch模型导出、onnx转pb模型、pb模型转saved_model等步骤，并解决了转换中可能遇到的问题，如版本兼容性和错误处理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Pytorch模型的TF-Serving部署

写在前面

目前PyTorch在学术界几乎完全盖过了Tensorflow，从个人体验来说：

PyTorch并没有表现出比Tensorflow性能差很多，基本是相当；

与Tensorflow的非Eager API相比，PyTorch上手、使用的难度相对来说非常低，并且调试非常容易（可断点可单步）；

以及笔者还没有体会到的动态图与静态图之争

但是在应用场景下仍然还有将模型使用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 = torch.tensor((), dtype=torch.int64)
input_ids = tensor.new_ones((1, 48)).to(device)
token_type_ids = tensor.new_ones((1, 48)).to(device)
attention_mask = tensor.new_ones((1, 48)).to(device)

# 关注点3
torch.onnx.export(model, 
                  (input_ids, attention_mask, token_type_ids),