Chuyển đổi mô hình PyTorch sang TensorFlow Lite: 4 phương pháp hiệu quả

Chuyển đổi mô hình học máy từ PyTorch sang TensorFlow Lite (TFLite) là một bước thiết yếu để triển khai mô hình trên thiết bị di động và nhúng. Tuy nhiên, quá trình này đòi hỏi sự hiểu biết sâu về cả hai framework và cách chúng biểu diễn các phép toán toán học. Dưới đây là 4 phương pháp được kiểm chứng thực tế, giúp đảm bảo chất lượng và hiệu năng sau khi chuyển đổi.

1. Sử dụng ONNX làm trung gian trung gian

Phương pháp này phân rã quá trình thành hai bước: PyTorch → ONNX → TensorFlow → TFLite. Đây là cách tiếp cận linh hoạt và phổ biến nhất, đặc biệt khi mô hình không được hỗ trợ trực tiếp bởi công cụ chuyển đổi chính thức của TensorFlow.

Bước 1: Xuất mô hình PyTorch sang định dạng ONNX

import torch
from torch.autograd import Variable

# Giả sử model đã được huấn luyện
dummy_input = torch.randn(1, 3, 224, 224)  # batch_size=1, channels=3, 224×224px

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=13,
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

Bước 2: Chuyển đổi ONNX sang TensorFlow SavedModel

pip install onnx-tf
python -m onnx_tf conversion -i model.onnx -o model_savedmodel/

Bước 3: Tạo tệp TFLite từ SavedModel

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("model_savedmodel")
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,  #サポートヒューリスティック
    tf.lite.OpsSet.SELECT_TF_OPS,    #miễn trừ TensorFlow op không có trong TFLite
]
tflite_model = converter.convert()

with open("model.tflite", "wb") as f:
    f.write(tflite_model)

2. Chuyển đổi trực tiếp từ PyTorch sang TFLite thông qua TensorFlow Keras

Phương pháp này phù hợp với các mô hình đơn giản hoặc có kiến trúc tương thích cao. Nó chuyển pytoch sang định dạng Keras trước, rồi sang TFLite.

# 1. Export PyTorch weights vào file
torch.save(model.state_dict(), "weights.pt")

# 2. Đọc và重建 mô hình trong TensorFlow/Keras
import tensorflow as tf

def build_keras_model_from_pytorch_weights():
    # Construct Keras model architecture matching the PyTorch version
    inputs = tf.keras.Input(shape=(3, 224, 224))
    x = tf.keras.layers.Conv2D(64, (3, 3), activation="relu")(inputs)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    outputs = tf.keras.layers.Dense(10, activation="softmax")(x)
    keras_model = tf.keras.Model(inputs, outputs)

    # Load PyTorch weights into Keras manually or use tools like `torch2keras`
    keras_model.load_weights("weights.h5")  # Must convert .pt → .h5 format
    return keras_model

# 3. Convert sang TFLite
keras_model = build_keras_model_from_pytorch_weights()
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

Bạn cần tự đảm bảo tương thích kiến trúc giữa hai framework. Ngoài ra, có thể dùng các thư viện như `torch2keras`, `polygraphy`, hoặc `mmdnn` để hỗ trợ chuyển đổi trọng số.

3. Dùng công cụ chuyển đổi chính thức của TensorFlow với PT2TF Bridge

Từ TensorFlow 2.12+, có tích hợp PyTorch-Signal (torch.fx) để xây dựng graph chuyển đổi nội tại. Việc này giúp giữ nguyên ngữ nghĩa phép toán cao hơn.

Trước tiên, import các module cần thiết từ `tensorflow.compiler.tf2torch`:

# Giả sử sử dụng phiên bản mới hỗ trợ FXGraph
from tensorflow.compiler.tf2torch import converter as tf2torch_converter

# Giả sử PyTorch model là torch.nn.Module
class MyModel(torch.nn.Module):
    def forward(self, x):
        return self.layer(x)

model = MyModel()
model.eval()

# Biên dịch GRAPH FX từereço PyTorch model
fx_g = torch.fx.symbolic_trace(model)

# Sử dụng converter mới
converter = tf2torch_converter.PyTorchToTensorFlowConverter()
tflite_model = converter.convert(fx_g, input_shapes=[(1, 3, 224, 224)])

with open("model_fx.tflite", "wb") as f:
    f.write(tflite_model)

Đây là phương pháp mới và đang trong giai đoạn thử nghiệm. Mức độ hỗ trợ còn hạn chế, nhưng hứa hẹn精度 cao hơn trong việc giữ nguyên mô hình gốc.

4. Tự xây dựng layer và import trọng số thủ công

Khi các phương pháp trên thất bại (ví dụ: mô hình có custom layers, custom ops, hoặc operation graph quá phức tạp), ta có thể tái hiện từng phần thủ công trong TFLite bằng cách:

  1. Xuất toàn bộ trọng số đã huấn luyện (thường là `state_dict`).
  2. Tạo lại toàn bộ kiến trúc bằng TensorFlow ops chuẩn.
  3. Gán trọng số đã export trước đó.
  4. Chuyển đổi sang TFLite như mô hình Keras thông thường.
# 1. Lưu trọng số dạng numpy array để dễ dùng trong TF
import numpy as np
weights_dict = {}
for name, param in model.named_parameters():
    weights_dict[name] = param.detach().cpu().numpy()

np.savez(" weights.npz", **weights_dict)

# 2. Rebuild dùng hàm lambda trong tf.data hoặc tf.function
class MyTFLiteModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.conv1 = tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='same')
        self.fc = tf.keras.layers.Dense(10)

    def call(self, x):
        x = self.conv1(x)
        x = tf.nn.relu(x)
        x = tf.reduce_mean(x, axis=[1, 2])
        return self.fc(x)

# 3. Assign weights thủ công từ npz file
reloaded_model = MyTFLiteModel()
reloaded_model.build(input_shape=(1, 3, 224, 224))

weights_data = np.load("weights.npz")
reloaded_model.conv1.set_weights([weights_data['conv1.weight'], weights_data['conv1.bias']])
reloaded_model.fc.set_weights([weights_data['fc.weight'].T, weights_data['fc.bias']])

# 4. Convert sang TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(reloaded_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

with open("quant_model.tflite", "wb") as f:
    f.write(tflite_quant_model)

Phương pháp này tốn nhiều công sức, nhưng đảm bảo khả năng kiểm soát tối đa — cực kỳ hữu ích trong các trường hợp nhạy cảm về hiệu quả hoặc bảo mật.

Đảm bảo chất lượng sau chuyển đổi

Sau khi chuyển đổi, bạn bắt buộc phải thực hiện các bước sau:

  • So sánh kết luận (inference output) giữa mô hình PyTorch gốc và TFLite (dùng cùng dataset nhỏ).
  • Đo lường số liệu latent: thời gian suy luận, RAM tiêu thụ, năng lượng.
  • Áp dụng Quantization nếu mục tiêu là giảm size model hoặc tăng tốc trên Edge TPU, NPU.
# Đo độ chính xác tương đối
def evaluate_tflite_vs_torch(tflite_path, test_loader):
    interpreter = tf.lite.Interpreter(model_path=tflite_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    correct = 0
    total = 0
    for images, labels in test_loader:
        interpreter.set_tensor(input_details[0]['index'], images.numpy())
        interpreter.invoke()
        tflite_output = interpreter.get_tensor(output_details[0]['index'])
        
        torch_output = model(images)
        
        pred_tflite = np.argmax(tflite_output, axis=1)
        pred_torch = torch.argmax(torch_output, axis=1).numpy()
        
        correct += (pred_tflite == pred_torch).sum()
        total += labels.size(0)

    print(f"Accuracy match: {100 * correct / total:.2f}%")

Lưu ý về_gradients_ và chất lượngbit

- giống như việc export sang ONNX, cần hiểu rõ về lưu đồ tính toán (computation graph), đặc biệt với các op như selfattention, group norm, hoặc batch norm — các op này có thể biểu diễn khác nhau giữa hai framework. - khi dùng quantization, hãy thử cả full integer (int8) và float16, đặc biệt với các mạng nhỏ hoặc khi độ chính xác ảnh hưởng mạnh đến trải nghiệm cuối người dùng. - một số mô hình sử dụng dynamic control flow (như loop, recursion) thường không hỗ trợ trong TFLite hiện tại — nên tái cấu trúc logic trước khi export.

Đăng vào ngày 24 tháng 7 lúc 18:14

Thẻ Phổ Biến

©2026 Thành phố Cuồng loạn