Chuyển đổi các kiểu máy PyTorch sang TF Lite

AI Edge Torch là một thư viện cho phép bạn chuyển đổi các mô hình PyTorch sang định dạng .tflite, nhờ đó, bạn có thể chạy các mô hình đó bằng TensorFlow Lite và MediaPipe. Điều này đặc biệt hữu ích đối với các nhà phát triển tạo ứng dụng di động chạy các mô hình hoàn toàn trên thiết bị. AI Edge Torch cung cấp phạm vi CPU rộng, với GPU ban đầu và hỗ trợ NPU.

Để bắt đầu chuyển đổi các mô hình PyTorch sang TF Lite, hãy sử dụng bắt đầu nhanh bộ chuyển đổi Pytorch. Để biết thêm thông tin, hãy xem kho lưu trữ GitHub của AI Edge Torch.

Nếu bạn đang chuyển đổi cụ thể Mô hình ngôn ngữ lớn (LLM) hoặc mô hình dựa trên bộ biến đổi, hãy sử dụng API Đèn pin tạo sinh. API này xử lý các chi tiết chuyển đổi dành riêng cho bộ biến đổi, chẳng hạn như tạo mô hình và lượng tử hoá.

Quy trình chuyển đổi

Các bước sau đây minh hoạ quy trình chuyển đổi toàn diện đơn giản từ mô hình PyTorch sang TensorFlow Lite.

Nhập đèn pin AI Edge

Bắt đầu bằng cách nhập gói pip AI Edge Torch (ai-edge-torch) cùng với PyTorch.

import ai_edge_torch
import torch

Đối với ví dụ này, chúng tôi cũng yêu cầu các gói sau:

import numpy
import torchvision

Khởi động và chuyển đổi mô hình

Chúng ta sẽ chuyển đổi ResNet18, một mô hình nhận dạng hình ảnh phổ biến.

resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()

Sử dụng phương thức convert từ thư viện AI Edge Torch để chuyển đổi mô hình PyTorch.

sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)

Sử dụng mô hình

Sau khi chuyển đổi mô hình Pytorch, bạn có thể chạy dự đoán bằng mô hình TF Lite mới đã chuyển đổi.

output = edge_model(*sample_inputs)

Bạn có thể xuất và lưu mô hình đã chuyển đổi ở định dạng .tflite để sử dụng trong tương lai.

edge_model.export('resnet.tflite')