AI

On-Device AI 프레임워크

_어니언_ 2024. 12. 5. 19:44

길게 설명할 필요 없이 요즘 온디바이스(on-device) AI 가 핫하다. 사이즈가 큰 모델을 서버에서 돌려 네트워크로 받아오는 것이 아닌, 로컬 디바이스에서 추론을 돌려 네트워크 비용과 보안 이슈를 걱정하지 않고 AI 를 적용하겠다는 것이다. 

 

이에 맞추어 모델을 최적화하고 경량화 하는 프레임워크들이 정말 우후죽순으로 나오고 있다. 대표적인 프레임워크를 살펴보면 다음과 같다.

 


 

ONNX (Open Neural Network Exchange)

  • 설명: 일반적인 모델 최적화 프레임워크. 다른 프레임워크 간 모델 호환성을 제공하는 표준이 된다.
  • 특징: PyTorch, TensorFlow 등의 프레임워크에서 모델을 변환하여 다양한 하드웨어에서 실행 가능. (단 기기 별 최적화 효율은 떨어질 수 있음)
  • 최적화 프레임 워크들 중 대장 정도라고 보면 된다.

TesorRT

  • 설명: NVIDIA CUDA 기반의 모델에 최적화
  • CUDA를 지원하는데 무슨 말이 더 필요한가?

OpenVINO

  • 설명: Intel 플랫폼에 최적화 된 AI 추론
  • 특징: CPU, GPU, NPU 등 다양한 모델 지원.
  • 활용: Intel 플랫폼 위에서 AI 추론

CoreML

  • 설명: Apple 생태계(iOS, macOS) 위에서 AI 모델 최적화를 위한 프레임워크
  • 특징: Metal API와 통합으로 GPU 가속, Swift 및 Objective-C와의 쉬운 통합
  • 활용: iOS 등 애플 기반 서비스에서 사용되는 AI

LiteRT (구 Tensorflow Lite)

  • 설명: Google TensorFlow 기반 모델을 경량화하는 Framework
  • 활용: TF 기반으로 모델을 학습하고 추론할 때 유용. 
  • MediaPipe 등 TFLite 기반으로 학습되고 최적화 된 on-device 모델 사례가 많아서 굉장히 유용함

SNPE (Qualcomm Snapdragon Neural Processing Engine)

  • 설명: Qualcomm Snapdragon 칩셋에 특화된 AI 최적화 프레임워크
  • 활용: 퀄컴 칩 기반 안드로이드 스마트폰에서 추론.

 


최적화 프레임워크 기본 사용법

위에 나열된 모든 AI 프레임워크는 기본적으로 다음과 같은 패턴을 가진다. 모델을 준비하고 Input 과 Output 을 지정해주면 프레임워크가 알아서 모델을 자체적으로 최적화 해주는 패턴이다. 간단하게 torch 에서 제공해주는 onnx 기능을 통해 살펴보자.

 

# export_model.py

import torch
import torch.onnx

# Define or load your PyTorch model
model = YourPyTorchModel()
model.eval()

# Dummy input for the model
dummy_input = torch.randn(1, 3, 224, 224)  # Adjust shape as needed for your model

# Export the model to ONNX format
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=11,  # Update the opset version as required
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"]
)

 

위 파이썬 스크립트를 보면 inference 에 사용될 pytorch 모델을 불러와 sample input 과 output 을 지정해서 모델에 넣어주면 ONNX 로 포팅된 모델을 저장할 수 있다. 

 

이후 ONNX 로 포팅된 모델로 다음과 같이 inference 를 돌릴 수 있다.

 

# inference.py

import onnxruntime
import numpy as np

# Load the ONNX model
onnx_model_path = "model.onnx"
session = onnxruntime.InferenceSession(onnx_model_path)

# Prepare input data
# Adjust input shape and data type to match the model's requirements
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)  # Example input

# Get input and output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Run inference
outputs = session.run([output_name], {input_name: input_data})

# Print the result
print("Inference result:", outputs[0])

 

위 코드를 보면 포팅된 모델을 `session` 으로 불러와 input_data 를 넣어주어 inference 된 결과를 받아오는 것을 볼 수 있다.

 

위는 단순 Torch 에 내장된 ONNX 기능이지만 이 외에 다른 프레임워크들 모두 위와 같은 패턴을 따른다고 볼 수 있다. ONNX 가 범용적인 하드웨어에서 실행할 수 있는 최적화 표준을 제공해준다면 각 플랫폼(Intel, Ap ple, NVIDIA 등) 에서 추론을 최적화하기 위해서는 플랫폼에 최적화 된 툴(OpenVINO, CoreML, TensorRT 등)로 inference 를 돌려야 한다. 생각보다 하드웨어들의 구성, 명령어 세트부터 가속화 지원 여부, 코어와 캐시 등 하드웨어 구성들이 모두 다르기에 최적의 성능을 내기 위해서는 각 플랫폼 별 최적화 된 프레임워크를 사용해야 한다.

 

그래서 혹시 AI 최적화를 공부해보고자 한다면 먼저 ONNX 로 모델을 변환해보고 변환된 모델을 TensorRT, OpenVINO 등 하드웨어에 최적화된 프레임워크로 로드해서 돌려보는 것을 추천한다.