4. PyTorch模型直接转换为TFLite

任务目标

将PyTorch训练好的模型转换为TFLite格式，以便在移动端或边缘设备上部署。本案例以经典CNN分类模型为例，走通 PyTorch → ONNX → TensorFlow → TFLite 的完整转换链路，并验证转换后的TFLite模型推理结果与原始PyTorch模型一致。

环境依赖

  • PyTorch ≥ 1.10（含 torch.onnx.export）
  • ONNX ≥ 1.12（onnx 与 onnx2tf）
  • TensorFlow ≥ 2.8（含 tf.lite.TFLiteConverter）
  • 安装：pip install onnx onnx2tf tensorflow

实现步骤与代码

核心思路：先用 torch.onnx.export 将PyTorch模型固化为ONNX，再用 onnx2tf 将ONNX转为TensorFlow SavedModel，最后用 TFLiteConverter 生成TFLite文件。以下代码一次完成全流程，并包含输入输出张量形状对齐与推理验证。

import torch
import torch.nn as nn
import numpy as np
import onnx
import tensorflow as tf
import os

# ---------- 1. 定义并训练一个极简PyTorch模型（演示用） ----------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 4, kernel_size=3, padding=1)
        self.fc = nn.Linear(4 * 8 * 8, 10)

    def forward(self, x):
        x = torch.relu(self.conv(x))
        x = torch.max_pool2d(x, 2)  # 8x8
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = SimpleCNN()
model.eval()

# 固定输入尺寸：batch=1, channel=1, H=W=16
dummy_input = torch.randn(1, 1, 16, 16)

# ---------- 2. PyTorch → ONNX ----------
onnx_path = "simple_cnn.onnx"
torch.onnx.export(
    model, dummy_input, onnx_path,
    input_names=["input"], output_names=["output"],
    dynamic_axes=None,  # 固定shape，便于TFLite转换
    opset_version=13
)

# 验证ONNX模型结构完整性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX export OK")

# ---------- 3. ONNX → TensorFlow SavedModel ----------
# onnx2tf会保留输入输出名，并自动处理NHWC/NCHW转换
saved_model_dir = "saved_model"
os.makedirs(saved_model_dir, exist_ok=True)
os.system(f"onnx2tf -i {onnx_path} -o {saved_model_dir} -nuo 1")

# ---------- 4. TensorFlow → TFLite ----------
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS  # 允许部分TF算子回退
]
tflite_model = converter.convert()
tflite_path = "simple_cnn.tflite"
with open(tflite_path, "wb") as f:
    f.write(tflite_model)
print("TFLite conversion OK")

# ---------- 5. 推理验证：对比PyTorch与TFLite输出 ----------
# 准备相同输入
test_input = torch.randn(1, 1, 16, 16)

# PyTorch推理
with torch.no_grad():
    pt_output = model(test_input).numpy()

# TFLite推理
interpreter = tf.lite.Interpreter(model_path=tflite_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 检查输入张量形状（TFLite通常为NHWC）
print("TFLite input shape:", input_details[0]['shape'])
tflite_input = test_input.numpy().transpose(0, 2, 3, 1)  # NCHW→NHWC
interpreter.set_tensor(input_details[0]['index'], tflite_input)
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])

# 比较输出（softmax前logits）
diff = np.abs(pt_output - tflite_output).max()
print(f"Max absolute difference: {diff:.6f}")
assert diff < 1e-4, "Output mismatch too large"
print("Verification PASSED: PyTorch and TFLite outputs match")

运行方式

  1. 将上述代码保存为 convert_pytorch_to_tflite.py，直接执行 python convert_pytorch_to_tflite.py。
  2. 脚本运行结束后，当前目录应生成 simple_cnn.onnx、saved_model/ 目录和 simple_cnn.tflite 三个产物。

验证方法与判定标准

  • 转换链路完整性：脚本依次打印 ONNX export OK 和 TFLite conversion OK，且无异常抛出。
  • 推理一致性：脚本最后输出 Max absolute difference 数值，并断言其小于 1e-4。若断言通过，打印 Verification PASSED，说明TFLite模型与原始PyTorch模型在相同输入下的输出误差在可接受范围内。
  • 产物可用性：可用 Netron 打开 simple_cnn.tflite 查看计算图结构，确认输入形状为 [1, 16, 16, 1]（NHWC），输出形状为 [1, 10]。

常见问题

  • 算子不支持：若转换时报错，可在 converter.target_spec.supported_ops 中保留 SELECT_TF_OPS（代码已包含），或尝试降低ONNX opset版本。
  • 动态维度：本案例固定了batch和输入尺寸。若需动态batch，需在 torch.onnx.export 中设置 dynamic_axes，并在TFLite转换时确认推理引擎支持动态shape。
  • 输入布局：PyTorch默认NCHW，TFLite通常使用NHWC。代码中已通过 transpose 手动对齐，若使用 onnx2tf 的默认参数，其内部已处理布局转换，此处显式转置仅为确保一致性。
← 上一章 📖 返回目录 下一章 →