跨越重重“障碍”，我从 PyTorch 转换为了 TensorFlow Lite 作者 | Ran Rubin 译者 | Sambodhi 策划 & 编辑 | 刘燕 本文作者分享了他在 PyTorch 到 TensorFlow 之间转换的经验，或许可以给我们一些启发。 本文最初发表在 Towards Data Science 博客，经原作者 Ran Rubin 授权，InfoQ 中文站翻译并分享。 简 介 我最近不得不将深度学习模型（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 import onnx import torch example_input = get_example_input() # exmample for the forward pass input pytorch_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 saved verbose=False, export_params=True, do_constant_folding=False, # fold constant values for optimization # do_constant_folding=True, # fold constant values for optimization input_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 机器。 from onnx_tf.backend import prepare import onnx TF_PATH = "./my_tf_model.pb" # where the representation of tensorflow model will be stored ONNX_PATH = "./my_model.onnx" # path to my existing ONNX model onnx_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 工具导出的内容是一个冻结的图，所以，这三个选项都帮不了我。 我在网上搜索了很久之后，这个家伙 基本上拯救了我。原来，TensorFlowv1是支持从冻结图进行转换的！我决定在剩下的代码中使用v1API。 在运行转换函数时，出现了一个奇怪的问 p 题，它与protobuf库有关。遵循 这个用户 的建议，我得以能够继续前进。 TF_PATH = "./my_tf_model.pb" # where the forzen graph is stored TFLITE_PATH = "./my_model.tflite" # protopuf needs your virtual environment to be explictly exported in the path os.environ["PATH"] = "/opt/miniconda3/envs/convert/bin:/opt/miniconda3/bin:/usr/local/sbin:...." # make a converter object from the saved tensorflow file converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(TF_PATH, # TensorFlow freezegraph .pb model file input_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 ops converter.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”，在 这位用户的评论 的帮助下，这个问题得到了解决。 最后，下面是用于测试的推理代码： import os import tensorflow as tf import numpy as np TFLITE_PATH = "./my_model.tflite" example_input = get_numpy_example() print(f"Using tensorflow {tf.__version__}") # make sure it's the nightly build os.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]('index'), example_input) interpreter.invoke() print(interpreter.get_tensor(output_details[0]('index'))) # printing the result 测试结果的平均误差为 2.66e-07。 希望我的经验对你们有用，祝好运。 作者介绍： Ran Rubin，DevOps-MLOps 工程师。着迷于将运维和机器学习世界结合在一起。 原文链接： https://towardsdatascience.com/my-journey-in-converting-pytorch-to-tensorflow-lite-d244376beed 活动推荐 未来十年是基础设施和应用场景的交替式变革？WICC 2020 全球互联网通信云大会即将重磅来袭，解读实时音视频直播的架构建设、国产化统一操作系统架构下的多边技术融合，点击【阅读原文】或扫描下方【二维码】，急速关注！ 今日荐文点击下方图片即可阅读 突发！美国出口限制清单新增 6 项新兴技术，芯片制造首当其冲 你也「在看」吗？?? 阅读原文