PyTorch 모델을 TF Lite로 변환

AI Edge Torch는 PyTorch 모델을 .tflite 형식으로 변환하여 TensorFlow Lite 및 MediaPipe로 이러한 모델을 실행할 수 있게 해주는 라이브러리입니다. 이 기능은 완전히 기기에서 모델을 실행하는 모바일 앱을 만드는 개발자에게 특히 유용합니다. AI Edge Torch는 초기 GPU 및 NPU 지원과 함께 광범위한 CPU 범위를 제공합니다.

PyTorch 모델을 TF Lite로 변환하려면 PyTorch 변환기 빠른 시작을 사용하세요. 자세한 내용은 AI Edge Torch GitHub 저장소를 참조하세요.

대규모 언어 모델 (LLM) 또는 변환 기반 모델을 특별히 변환하는 경우 모델 작성 및 양자화와 같은 변환기별 변환 세부정보를 처리하는 Generative Torch API를 사용합니다.

전환 워크플로

다음 단계에서는 PyTorch 모델을 TensorFlow Lite로 간단하게 엔드 투 엔드로 변환하는 방법을 보여줍니다.

AI Edge Torch 가져오기

먼저 AI Edge Torch (ai-edge-torch) pip 패키지를 PyTorch와 함께 가져옵니다.

import ai_edge_torch
import torch

이 예에서는 다음 패키지도 필요합니다.

import numpy
import torchvision

모델 초기화 및 변환

여기서는 인기 있는 영상 인식 모델인 ResNet18을 변환합니다.

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

AI Edge Torch 라이브러리의 convert 메서드를 사용하여 PyTorch 모델을 변환합니다.

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

모델 사용

Pytorch 모델을 변환한 후 새롭게 변환된 TF Lite 모델로 추론을 실행할 수 있습니다.

output = edge_model(*sample_inputs)

나중에 사용할 수 있도록 변환된 모델을 .tflite 형식으로 내보내고 저장할 수 있습니다.

edge_model.export('resnet.tflite')