如何在TensorFlow中加载MLIR格式的ResNet50模型并推理？

方案一：将MLIR转回TFLite后使用TFLite解释器推理 #

这是最直接的落地路径，利用反向转换回到TFLite格式后，使用官方成熟的TFLite推理API完成预测，适配你原有的预处理/后处理逻辑。

步骤1：MLIR转回TFLite FlatBuffer #

通过flatbuffer_translate工具执行反向转换：

bazel run //tensorflow/compiler/mlir/lite:flatbuffer_translate -- --mlir-to-tflite-flatbuffer resnet50.mlir -o resnet50_reconverted.tflite

步骤2：Python中加载TFLite模型并推理 #

替换原Keras模型的推理逻辑，改用TFLite解释器实现预测：

import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
import tensorflow as tf

# 图像预处理（与原代码完全一致）
img = image.load_img("Panda.JPG", target_size=(224, 224))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

# 加载TFLite模型并初始化
interpreter = tf.lite.Interpreter(model_path="resnet50_reconverted.tflite")
interpreter.allocate_tensors()

# 获取输入输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 传入输入数据并执行推理
interpreter.set_tensor(input_details[0]['index'], x)
interpreter.invoke()

# 获取推理结果并解析
predictions = interpreter.get_tensor(output_details[0]['index'])
predicted_classes = decode_predictions(predictions, top=9)
print(predicted_classes)

方案二：直接使用TensorFlow MLIR Python API加载执行 #

如果需要直接操作MLIR中间表示（比如自定义优化、调试场景），可以用TensorFlow的MLIR Python绑定编译并运行MLIR模块，不过该路径对MLIR规范的适配要求较高。

Python代码实现 #

import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
import tensorflow as tf
from tensorflow.compiler.mlir.tensorflow import tf_saved_model

# 图像预处理（与原代码一致）
img = image.load_img("Panda.JPG", target_size=(224, 224))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

# 读取MLIR文件内容
with open("resnet50.mlir", "r") as f:
    mlir_text = f.read()

# 将TFLite风格MLIR转换为TensorFlow SavedModel规范的MLIR
context = tf.mlir.Context()
module = tf.mlir.parse_string(mlir_text, context=context)
tf_saved_model.convert(module, saved_model_dir="temp_resnet_savedmodel")

# 加载转换后的SavedModel并执行推理
loaded_model = tf.saved_model.load("temp_resnet_savedmodel")
infer_func = loaded_model.signatures["serving_default"]

# 运行推理并解析结果
predictions = infer_func(tf.convert_to_tensor(x))[list(infer_func.structured_outputs.keys())[0]].numpy()
predicted_classes = decode_predictions(predictions, top=9)
print(predicted_classes)

注意事项 #

• 方案二要求MLIR模块能兼容TensorFlow SavedModel规范，若你的MLIR是纯TFLite专属格式，方案一的转换路径会更稳定可靠。
• 直接操作MLIR需要对MLIR语法、TensorFlow的MLIR转换流程有基础了解，否则易出现格式不兼容的报错。

内容的提问来源于stack exchange，提问作者Elouan Jeannot