.tflite — bản chất, cấu trúc và ứng dụng của định dạng mô hình TFLite

Định dạng .tflite là gì

.tflite là biểu diễn tuần tự hóa của mô hình học máy được tối ưu hóa cho suy luận trên các thiết bị có bộ nhớ và sức mạnh tính toán hạn chế. Không giống như định dạng SavedModel lưu trữ đồ thị trong protobuf và yêu cầu tài nguyên đáng kể để tải, .tflite sử dụng FlatBuffers — một thư viện tuần tự hóa với quyền truy cập dữ liệu không sao chép.

Định dạng .tflite được thiết kế có tính đến đặc thù của nền tảng di động: tiêu thụ bộ nhớ tối thiểu khi tải, khởi động nhanh và hỗ trợ trọng số lượng tử hóa. Tệp .tflite không hỗ trợ huấn luyện — chỉ suy luận. Đây là sự khác biệt cơ bản so với các định dạng TensorFlow đầy đủ, cho phép giảm đáng kể kích thước thời gian chạy.

Theo Google Research, 2024, các mô hình ở định dạng .tflite với lượng tử hóa INT8 tải nhanh hơn 60% so với các mô hình FP32 tương đương và tiêu thụ ít hơn 40% RAM trong quá trình suy luận.

Cấu trúc bên trong của tệp .tflite

Tệp .tflite bao gồm nhiều phần được tổ chức theo lược đồ FlatBuffers. Phần chính là Model, chứa đồ thị tính toán (SubGraph), danh sách toán tử (OperatorCodes) và bộ đệm trọng số. Mỗi SubGraph chứa các tensor, toán tử và chỉ mục đầu vào/đầu ra.

Phần OperatorCodes chứa định danh của tất cả các toán tử được sử dụng trong mô hình — Conv2D, DepthwiseConv2D, FullyConnected, Softmax và các toán tử khác. Mỗi toán tử có một mã từ danh sách BuiltinOperator được xác định trước. Nếu mô hình chứa các toán tử không có trong danh sách, chúng được đánh dấu là Custom và yêu cầu triển khai thông qua delegate hoặc toán tử tùy chỉnh.

Phần Buffers chứa dữ liệu nhị phân của trọng số mô hình. Tùy thuộc vào chế độ lượng tử hóa, trọng số có thể được lưu trữ ở FP32, FP16, INT8 hoặc ở định dạng lượng tử hóa với các tham số tỷ lệ. Các bộ đệm được ánh xạ trực tiếp vào bộ nhớ thông qua mmap, loại bỏ việc sao chép thêm.

Các phần chính của tệp .tflite

Lượng tử hóa và tối ưu hóa định dạng

.tflite hỗ trợ ba chế độ lượng tử hóa. Lượng tử hóa số nguyên hoàn toàn INT8 — trọng số và kích hoạt được biểu diễn dưới dạng số nguyên 8 bit. Việc chuyển đổi yêu cầu một tập dữ liệu đại diện để hiệu chỉnh phạm vi kích hoạt. Kích thước mô hình giảm 4 lần.

Lượng tử hóa FP16 — trọng số được chuyển đổi thành số dấu phẩy động 16 bit. Chế độ này không yêu cầu hiệu chỉnh và giảm thiểu mất độ chính xác. Các kích hoạt vẫn ở FP32. Kích thước mô hình giảm 2 lần. Được khuyến nghị cho các thiết bị hỗ trợ FP16 trên GPU và NPU.

Lượng tử hóa động — trọng số được chuyển đổi thành INT8 nhưng các kích hoạt được tính toán trong FP32. Kích thước mô hình giảm 4 lần trong khi độ chính xác vẫn cao hơn so với INT8 hoàn toàn. Theo Google ML Performance Benchmarks, chế độ này là tối ưu cho các mô hình NLP và các mô hình có độ nhạy cao về độ chính xác.

Tạo .tflite từ mô hình TensorFlow

Để chuyển đổi mô hình sang .tflite, hãy sử dụng TFLite Converter — một thành phần TensorFlow có sẵn thông qua API Python tf.lite.TFLiteConverter. Bộ chuyển đổi hỗ trợ ba nguồn: SavedModel, Keras H5 và ConcreteFunction. Một ví dụ tối thiểu về chuyển đổi mô hình Keras:

import tensorflow as tf

# Tải mô hình Keras đã huấn luyện từ tệp
model = tf.keras.models.load_model("my_model.h5")

# Chuyển đổi sang .tflite với lượng tử hóa động
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# Ghi mô hình đã chuyển đổi vào tệp .tflite
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

Để chuyển đổi với lượng tử hóa INT8 hoàn toàn, cần có tập dữ liệu hiệu chỉnh. Cung cấp representative_dataset để xác định phạm vi kích hoạt:

converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
tflite_quant_model = converter.convert()

Công cụ làm việc với .tflite

Google cung cấp một bộ công cụ để phân tích và tối ưu hóa các mô hình .tflite. TFLite Benchmark Tool đo thời gian suy luận trên thiết bị và hiển thị thống kê cho từng toán tử. TFLite Model Inspector trực quan hóa đồ thị tính toán và hiển thị kích thước tensor.

Netron là công cụ trực quan hóa mô hình đa nền tảng hỗ trợ .tflite cùng với ONNX, Core ML và PyTorch. Netron hiển thị cấu trúc đồ thị, tham số của từng toán tử và kết nối giữa các tensor. Hữu ích cho việc gỡ lỗi trong quá trình chuyển đổi.

TFLite Metadata API cho phép thêm siêu dữ liệu vào tệp .tflite: mô tả mô hình, định dạng dữ liệu đầu vào, đơn vị đo lường, thông tin tác giả. Siêu dữ liệu được Task Library sử dụng để tự động cấu hình tiền xử lý và hậu xử lý.

Ví dụ tải và chạy .tflite

Ví dụ tải mô hình .tflite trên iOS với Swift API. Swift API của TFLite cung cấp Interpreter để tải mô hình từ gói và chạy suy luận. Chỉ định delegate để tăng tốc phần cứng thông qua Core ML Delegate:

import TensorFlowLite

guard let modelPath = Bundle.main.path(
    forResource: "model", ofType: "tflite"
) else { return }

let interpreter = try Interpreter.init(modelPath: modelPath)
try interpreter.allocateTensors()

// Chuẩn bị dữ liệu đầu vào cho mô hình
let inputData = Data.init(count: 1 * 224 * 224 * 3)
try interpreter.copy(inputData, toInputAt: 0)

// Chạy suy luận mô hình
try interpreter.invoke()

// Đọc dữ liệu tensor đầu ra
let outputTensor = try interpreter.output(at: 0)
let results = outputTensor.data.withUnsafeBytes {
    Array($0.bindMemory(to: Float32.self))
}