pth转onnx转tflite并固定输入（附代码和实验结果）

文章目录

• 前言
• 一、pth2onnx
• 二、onnx2tf
• • 2.1安装依赖包
  • 2.2 转换代码
  • 2.3命令解释
• 三、转换测试
• • 3.1 onnx输入输出：
  • tflite_float32输入输出：
  • 3.3 tflite_in8输入输出：
  • 3.4 大小比较
  • 3.5超分测试
  • • 原图：
    • 3.5.1 onnx：
    • 3.5.2 tflite_float32：
    • 3.5.2 tflite_int8：
• 总结

前言

最近工程上在做基于端侧AI的超分模型，模型量化要从torch转换成onnx，再转为tflite。为了分别测试，我们固定输入输出和batch批次

一、pth2onnx

pth模型转onnx模型采用torch框架现有转换函数：torch.onnx.export（），函数使用过程中注意要设置输入输出name，以及将batch维度设置为非动态，以便后续onnx转tflite。
具体转换代码如下：

import torch
import torch.onnx
from models.mobilesr import MobileSR
import os

model_path = os.path.join('./mobilesr_x4.pth')
device = torch.device('cpu')
model = MobileSR()
model.load_state_dict(torch.load(model_path,map_location=device)['net'], strict=True)
model.eval()

device = torch.device('cpu')
input_tensor = torch.randn(1, 3, 50, 50, device=device)

torch.onnx.export(model, input_tensor, "../model_zoo/mobilesr_1017_NHWC.onnx",
                 keep_initializers_as_inputs=False, verbose=True,
                 input_names=["input"], output_names=["output"], export_params=True, opset_version=15)

print('----------Down!!!----------Down!!!-----------')

先转换成(N,C,H,W)，然后在onnx2tf再加上transpose转换成(N,H,W,C)

二、onnx2tf

2.1安装依赖包

pip install onnx2tf
pip install nvidia-pyindex
pip install sng4onnx
pip install onnx-graphsurgeon

2.2 转换代码

#默认生成float32和float16
onnx2tf -i mobilesr_1017_NHWC.onnx#会在当前文件夹生成saved_model，保存saved model和tflite格式文件
onnx2tf -i mobilesr_1017_NHWC.onnx -oiqt -qt per-channel -iqd int8 -oqd int8#

2.3命令解释

这里是命令中每个参数的含义：

• -i mobilesr_1017_NHWC.onnx：指定输入的 ONNX 模型文件路径。
• -oiqt：指定输出整数量化的 TensorFlow Lite 模型。
• -qt per-channel：指定量化类型为“每通道（per-channel）”量化。你也可以选择“每张量（per-tensor）”量化。
• -iqd int8：指定输入量化的数据类型为 INT8。
• -oqd int8：指定输出量化的数据类型为 INT8。
• 这个命令会将 mobilesr_1017_NHWC.onnx 模型转换为一个 INT8 量化的 TensorFlow Lite 模型，并将结果保存在默认的输出文件夹 saved_model 中。如果你想要指定一个不同的输出文件夹，可以使用 -o 参数，例如 -o /path/to/output/folder。

三、转换测试

3.1 onnx输入输出：

tflite_float32输入输出：

3.3 tflite_in8输入输出：

3.4 大小比较

onnx转换成tflite文件大小减少了23kb， tflite_int8比tflite_fl32大小减少了一倍多

3.5超分测试

原图：

3.5.1 onnx：

3.5.2 tflite_float32：

3.5.2 tflite_int8：

测试代码：
1.onnx
代码如下（示例）：

import numpy as np
import cv2
import onnxruntime as ort

# 步骤1：读取图片
image = cv2.imread('test.jpg')

# 步骤2：调整大小到模型输入尺寸
# 确保输入图片大小为(50, 50)
resized_image = cv2.resize(image, (50, 50))

# 步骤3：转换维度并归一化
# 将图片从[h, w, c]转换为[1, c, h, w]并转换为FLOAT32
resized_image = np.transpose(resized_image, (2, 0, 1)).astype(np.float32)
resized_image = np.expand_dims(resized_image, axis=0) / 255.0  # 归一化到[0, 1]

# 步骤4：使用模型进行推理
session = ort.InferenceSession('mobilesr_1017_NHWC.onnx')
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 运行模型推理
outputs = session.run([output_name], {input_name: resized_image})

# 步骤5：移除批次维度并反归一化
output_image = outputs[0][0]  # 移除批次维度
output_image = (output_image * 255).astype(np.uint8)  # 反归一化到[0, 255]

# 步骤6：调整输出图片形状到(200, 200, 3)
output_image = np.transpose(output_image, (1, 2, 0))  # 将形状从 [3, 200, 200] 转换为 [200, 200, 3]

# 步骤7：保存结果
cv2.imwrite('result_onnx.jpg', output_image)

2.tflite_float32

import cv2
import tensorflow as tf

# 步骤1：读取图片
image = cv2.imread('test.jpg')

# 步骤2：调整大小到模型输入尺寸
# 确保输入图片大小为(50, 50)
resized_image = cv2.resize(image, (50, 50))

# 步骤3：转换维度并归一化
# 将图片从[h, w, c]转换为[1, h, w, c]并转换为FLOAT32
resized_image = np.expand_dims(resized_image, axis=0).astype(np.float32) / 255.0  # 归一化到[0, 1]

# 步骤4：使用模型进行推理
interpreter = tf.lite.Interpreter(model_path='mobilesr_1017_NHWC_float32.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 设置输入张量
interpreter.set_tensor(input_details[0]['index'], resized_image)

interpreter.invoke()

# 获取输出张量
output_tensor = interpreter.get_tensor(output_details[0]['index'])

# 步骤5：移除批次维度
output_image = output_tensor[0]

# 步骤6：调整输出图片大小到(200, 200)
# 由于模型输出应该是超分辨率后的图片，我们需要将其调整到(200, 200)
output_image = cv2.resize(output_image, (200, 200))

# 步骤7：将输出图像转换为UINT8并保存结果
output_image = (output_image * 255).astype(np.uint8)  # 反归一化到[0, 255]
cv2.imwrite('result_fl32.png', output_image)

3.tflite_int8

import numpy as np
import cv2
import tensorflow as tf

# 步骤1：读取图片
image = cv2.imread('test.jpg')

# 步骤2：调整大小到模型输入尺寸
# 确保输入图片大小为(56, 56)，因为模型输入是 (None, 56, 56, 3)
resized_image = cv2.resize(image, (50, 50))

# 步骤3：转换维度并归一化
# 将图片从[h, w, c]转换为[1, h, w, c]并转换为INT8
resized_image = np.expand_dims(resized_image, axis=0).astype(np.float32) / 255.0  # 归一化到[0, 1]

# 步骤4：使用模型进行推理
interpreter = tf.lite.Interpreter(model_path='mobilesr_1017_NHWC_full_integer_quant.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 设置输入张量
# 由于输入是INT8，我们需要将像素值从[0, 1]缩放到[-128, 127]的范围
resized_image_int8 = (resized_image * 127.0).astype(np.int8)
interpreter.set_tensor(input_details[0]['index'], resized_image_int8)

interpreter.invoke()

# 获取输出张量
output_tensor = interpreter.get_tensor(output_details[0]['index'])

# 步骤5：移除批次维度
output_image = output_tensor[0]

# 步骤6：调整输出图片大小到(200, 200)
# 由于模型输出应该是超分辨率后的图片，我们需要将其调整到(200, 200)
output_image = cv2.resize(output_image, (200, 200))

# 步骤7：将输出图像转换为UINT8并保存结果
# 反归一化到[0, 255]的范围
output_image = (output_image.astype(np.float32) / 127.0 * 255).astype(np.uint8)
cv2.imwrite('result_int8.png', output_image)import numpy as np

总结

模型超分的结果不是非常理想，原因还在排查中，可能是因为图片输入太小，也可能因为测试代码没有处理好输入输出。总体上量化为tflite_float32，是目前最快和最优的超分模型。