首页
学习
活动
专区
工具
TVP
发布
精选内容/技术社群/优惠产品,尽在小程序
立即前往

跨越重重“障碍”,我从 PyTorch 转换为了 TensorFlow Lite

本文最初发表在 Towards Data Science 博客,经原作者 Ran Rubin 授权,InfoQ 中文站翻译并分享。

本文作者分享了他在 PyTorch 到 TensorFlow 之间转换的经验,或许可以给我们一些启发。

简 介

我最近不得不将深度学习模型(MobileNetV2的变体)从PyTorch转换为TensorFlow Lite。这是一个漫长而复杂的旅程。需要跨越很多障碍才能成功。我发现自己从 StackOverflow 帖子和 GitHub 的问题中搜集了一些信息。我的目标是分享我的经验,以帮助其他像我一样“迷失”的人。

免责声明:本文并非关于如何正确进行转换的指南。我只想分享我的经验,但我也有可能做错(尤其是因为我没有 TensorFlow 的经验)。

任 务

将深度学习模型(MobileNetV2 变体)从 PyTorch 转换为 TensorFlow Lite,转换过程应该是这样的:

PyTorch → ONNX → TensorFlow → TFLite

测 试

为了测试转换后的模型,我生成了一组大约 1000 个输入张量,并为每个模型计算了 PyTorch 模型的输出。这个集合后来被用来测试每个转换后的模型,方法是通过一个平均误差度量,在整个集合中将它们的输出与原始输出进行比较。在相同的输入下,平均误差反映了在相同的输入下,转换后的模型输出与原始 PyTorch 模型输出相比有多大的不同。

我决定将平均误差小于 1e-6 的模型视为成功转换的模型。

可能还需要注意的是,我在张量中添加了批维度,尽管它为 1。我没有理由这么做,除了来自我以前将 PyTorch 转换为 DLC 模型的经验的直觉。

将 PyTorch 转换为 ONNX

这绝对是最简单的部分。这主要归功于 PyTorch 的优秀文档,例如 TORCH.ONNX 的文档和《(可选)将模型从 PyTorch 导出到 ONNX 并使用 ONNX 运行时运行》( (Optional) Exporting a model from pytorch to onnx and running it using onnx runtime)。

要求:

  • ONNX == 1.7.0
  • PyTorch == 1.5.1

代码语言:javascript
复制
import onnximport torchexample_input = get_example_input() # exmample for the forward pass inputpytorch_model = get_pytorch_model()ONNX_PATH="./my_model.onnx"torch.onnx.export(model=pytorch_model,args=example_input,f=ONNX_PATH, # where should it be savedverbose=False,export_params=True,do_constant_folding=False,  # fold constant values for optimization# do_constant_folding=True,   # fold constant values for optimizationinput_names=['input'],output_names=['output'])onnx_model = onnx.load(ONNX_PATH)onnx.checker.check_model(onnx_model)

Python 到 ONNX 的转换

新创建的 ONNX 模型在我的示例输入上进行了测试,得到的平均误差为 1.39e-06。

请注意,你必须将torch.tensor示例转换为它们的等效np.array,才能通过 ONNX 模型运行它。

将 ONNX 转换到 TensorFlow

现在,我有了 ONNX 模型,为了转换成 TensorFlow,我使用了 ONNX-TensorFlow v1.6.0 )库。我并没有使用 TensorFlow 的经验,所以我知道这是事情变得有挑战性的地方。

要求:

  • TensorFlow == 2.2.0(这是 onnx-tensorflow 的先决条件。不过,它也适用于 tf-nightly 版本2.4.0-dev20200923)。
  • tensorflow-addons == 0.11.2
  • onnx-tensorflow == 1.6.0

我也不知道为什么,但这种转换只能用在我的 GPU 机器。

代码语言:javascript
复制
from onnx_tf.backend import prepareimport onnxTF_PATH = "./my_tf_model.pb" # where the representation of tensorflow model will be storedONNX_PATH = "./my_model.onnx" # path to my existing ONNX modelonnx_model = onnx.load(ONNX_PATH)  # load onnx model# prepare function converts an ONNX model to an internel representation# of the computational graph called TensorflowRep and returns# the converted representation.tf_rep = prepare(onnx_model)  # creating TensorflowRep object# export_graph function obtains the graph proto corresponding to the ONNX# model associated with the backend representation and serializes# to a protobuf file.tf_rep.export_graph(TF_PATH)

ONNX 到 TensorFlow 的转换 我在创建的 对象运行了测试(这里是使用它进行推理的示例)。运行超级慢(大约有 1 小时,而不是几秒钟!),所以这让我很担心。然而,最终测试的平均误差为 6.29e-07,所以我决定继续。

此时最大的问题是——它导出了什么?这个.pb文件又是什么?

我在网上搜索一番后,才意识到这是 tf.Graph的一个实例。现在剩下要做的就是把它转换成 TensorFlow Lite。

将 TensorFlow 转换到 TensorFlow Lite

这就是事情对我来说非常棘手的地方。据我所知, TensorFlow 提供了 3 种方法来将 TF 转换为 TFLite :SavedModel、Keras 和具体函数。可是我不太熟悉这些选项,但我已经知道 onnx-tensorflow 工具导出的内容是一个冻结的图,所以,这三个选项都帮不了我。

我在网上搜索了很久之后,这个家伙基本上拯救了我。原来,TensorFlow v1是支持从冻结图进行转换的!我决定在剩下的代码中使用v1API。

在运行转换函数时,出现了一个奇怪的问 p 题,它与protobuf库有关。遵循这个用户的建议,我得以能够继续前进。

代码语言:javascript
复制
 TF_PATH = "./my_tf_model.pb" # where the forzen graph is storedTFLITE_PATH = "./my_model.tflite"# protopuf needs your virtual environment to be explictly exported in the pathos.environ["PATH"] = "/opt/miniconda3/envs/convert/bin:/opt/miniconda3/bin:/usr/local/sbin:...."# make a converter object from the saved tensorflow fileconverter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(TF_PATH,  # TensorFlow freezegraph .pb model fileinput_arrays=['input'], # name of input arrays as defined in torch.onnx.export function before.output_arrays=['output'] # name of output arrays defined in torch.onnx.export function before.)# tell converter which type of optimization techniques to use# to view the best option for optimization read documentation of tflite about optimization# go to this link https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional# converter.optimizations = [tf.compat.v1.lite.Optimize.DEFAULT]converter.experimental_new_converter = True# I had to explicitly state the opsconverter.target_spec.supported_ops = [tf.compat.v1.lite.OpsSet.TFLITE_BUILTINS,tf.compat.v1.lite.OpsSet.SELECT_TF_OPS]tf_lite_model = converter.convert()# Save the model.with open(TFLITE_PATH, 'wb') as f:f.write(tf_lite_model)

TF 冻结图到 TFLite 你可能会认为,在经历了所有这些麻烦之后,在新创建的tflite模型上运行推理可以平静地进行。但是,我的麻烦并没有就此结束,更多的问题出现了。

其中之一与名为“ ops ”的东西有关(一个带有“Ops that can be supported by the Flex.”的错误消息)。经过一番搜索,我才意识到,我的模型架构需要在转换之前显式地启用一些操作符(见上文)。

然后,我发现我的网络使用的许多操作仍在开发中,因此正在运行的 TensorFlow 版本 2.2.0 无法识别它们。通过安装 TensorFlow 的 nightly 版本(特别是nightly==2.4.0.dev20299923),才解决了这一问题。

我遇到的另一个错误是“The Conv2D op currently only supports the NHWC tensor format on the CPU. The op was given the format: NCHW”,在这位用户的评论的帮助下,这个问题得到了解决。

最后,下面是用于测试的推理代码:

代码语言:javascript
复制
 import osimport tensorflow as tfimport numpy as npTFLITE_PATH = "./my_model.tflite"example_input = get_numpy_example()print(f"Using tensorflow {tf.__version__}") # make sure it's the nightly buildos.environ["CUDA_VISIBLE_DEVICES"] = "-1"interpreter = tf.compat.v1.lite.Interpreter(model_path=TFLITE_PATH)interpreter.allocate_tensors()input_details = interpreter.get_input_details()output_details = interpreter.get_output_details()interpreter.set_tensor(input_details 0 , example_input)interpreter.invoke()print(interpreter.get_tensor(output_details 0 )) # printing the result

测试结果的平均误差为 2.66e-07。

希望我的经验对你们有用,祝你们好运。

作者介绍:

Ran Rubin,DevOps-MLOps 工程师。着迷于将运维和机器学习世界结合在一起。

原文链接:

https://towardsdatascience.com/my-journey-in-converting-pytorch-to-tensorflow-lite-d244376beed

  • 发表于:
  • 本文为 InfoQ 中文站特供稿件
  • 首发地址https://www.infoq.cn/article/ACQFj4IqyKL48NEXoeox
  • 如有侵权,请联系 cloudcommunity@tencent.com 删除。

扫码

添加站长 进交流群

领取专属 10元无门槛券

私享最新 技术干货

扫码加入开发者社群
领券