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

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)

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)

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)

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[]('index'), example_input)interpreter.invoke()print(interpreter.get_tensor(output_details[]('index'))) # printing the result