Chuyển đổi mô hình PyTorch sang LiteRT

AI Edge Torch là một thư viện cho phép bạn chuyển đổi các mô hình PyTorch thành định dạng .tflite, cho phép bạn chạy các mô hình đó bằng LiteRT và MediaPipe. Điều này đặc biệt hữu ích cho 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 khả năng hỗ trợ GPU và NPU ban đầu.

Để bắt đầu chuyển đổi mô hình PyTorch sang LiteRT, hãy sử dụng hướng dẫn nhanh về trình 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ể các Mô hình ngôn ngữ lớn (LLM) hoặc mô hình dựa trên trình chuyển đổi, hãy sử dụng API Torch tạo sinh. API này xử lý các chi tiết chuyển đổi dành riêng cho trình chuyển đổi, chẳng hạn như việc tạo mô hình và lượng tử hoá.

Quy trình chuyển đổi

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

Nhập AI Edge Torch

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 ta cũng cần có các gói sau:

import numpy
import torchvision

Khởi tạo 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 quy trình suy luận bằng mô hình LiteRT mới được 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')