클라우드를 끊어도 AI가 돌아간다 — Edge AI 온디바이스 배포 파이프라인 구현

import tensorflow as tf
 
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model/")
 
# Post-Training Quantization (PTQ) 적용
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.int8]
 
# 대표 데이터셋으로 범위 보정 (INT8 정적 양자화)
def representative_dataset():
    for data in calibration_data:
        yield [data.astype("float32")]
 
converter.representative_dataset = representative_dataset
tflite_model = converter.convert()
 
with open("model_int8.tflite", "wb") as f:
    f.write(tflite_model)

import torch
import torch.nn.functional as F
 
def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    # KL divergence: 두 확률 분포 사이의 거리를 측정하는 함수입니다.
    # 교사 모델의 소프트 확률 분포를 학생이 모방하도록 유도합니다.
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / temperature, dim=1),
        F.softmax(teacher_logits / temperature, dim=1),
        reduction="batchmean"
    ) * (temperature ** 2)
 
    # 하드 타겟 손실: 실제 레이블로 학습
    hard_loss = F.cross_entropy(student_logits, labels)
 
    return alpha * soft_loss + (1 - alpha) * hard_loss

import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.nnapi.NnApiDelegate
import android.content.res.AssetManager
import android.graphics.Bitmap
import java.io.FileInputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.MappedByteBuffer
import java.nio.channels.FileChannel
 
class EdgeInferenceEngine(private val assetManager: AssetManager, modelFileName: String) {
 
    private val nnApiDelegate = NnApiDelegate()
    private val interpreter: Interpreter
    private val inputSize = 224
    private val numClasses = 1000
 
    init {
        val options = Interpreter.Options().apply {
            addDelegate(nnApiDelegate)
            // NPU 폴백 시 CPU 멀티스레드로 처리
            setNumThreads(4)
        }
        interpreter = Interpreter(loadModelFile(modelFileName), options)
    }
 
    fun classify(bitmap: Bitmap): FloatArray {
        val inputBuffer = preprocessImage(bitmap)
        val outputBuffer = Array(1) { FloatArray(numClasses) }
        interpreter.run(inputBuffer, outputBuffer)
        return outputBuffer[0]
    }
 
    // NnApiDelegate는 AutoCloseable을 구현하므로 반드시 해제해야 합니다.
    // Activity나 Fragment의 onDestroy()에서 호출하시면 됩니다.
    fun close() {
        interpreter.close()
        nnApiDelegate.close()
    }
 
    private fun loadModelFile(fileName: String): MappedByteBuffer {
        val fileDescriptor = assetManager.openFd(fileName)
        val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
        return inputStream.channel.map(
            FileChannel.MapMode.READ_ONLY,
            fileDescriptor.startOffset,
            fileDescriptor.declaredLength
        )
    }
 
    private fun preprocessImage(bitmap: Bitmap): ByteBuffer {
        val resized = Bitmap.createScaledBitmap(bitmap, inputSize, inputSize, true)
        // allocateDirect: JVM 힙 외부에 메모리를 할당해 JNI 복사 비용을 없앱니다.
        // nativeOrder(): 디바이스 CPU의 바이트 순서(엔디안)에 맞춥니다.
        val buffer = ByteBuffer.allocateDirect(4 * 3 * inputSize * inputSize)
        buffer.order(ByteOrder.nativeOrder())
 
        val pixels = IntArray(inputSize * inputSize)
        resized.getPixels(pixels, 0, inputSize, 0, 0, inputSize, inputSize)
 
        for (pixel in pixels) {
            // 훈련 시 사용한 ImageNet 정규화 파라미터와 반드시 일치해야 정확도가 나옵니다.
            buffer.putFloat(((pixel shr 16 and 0xFF) / 255.0f - 0.485f) / 0.229f)
            buffer.putFloat(((pixel shr 8 and 0xFF) / 255.0f - 0.456f) / 0.224f)
            buffer.putFloat(((pixel and 0xFF) / 255.0f - 0.406f) / 0.225f)
        }
        return buffer
    }
}

from typing import Any
import onnxruntime as ort
import numpy as np
from PIL import Image
 
def create_session(model_path: str) -> ort.InferenceSession:
    # 사용 가능한 프로바이더 자동 선택 (CUDA > CoreML > CPU 순)
    providers = ort.get_available_providers()
    print(f"사용 가능한 실행 프로바이더: {providers}")
 
    session_options = ort.SessionOptions()
    session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    session_options.intra_op_num_threads = 4
 
    return ort.InferenceSession(model_path, session_options, providers=providers)
 
def run_inference(session: ort.InferenceSession, image_path: str) -> np.ndarray:
    # convert("RGB")를 명시하지 않으면 RGBA PNG나 흑백 이미지에서 shape mismatch가 납니다.
    img = Image.open(image_path).convert("RGB").resize((224, 224))
    input_data = np.array(img, dtype=np.float32) / 255.0
    input_data = (input_data - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
    input_data = np.transpose(input_data, (2, 0, 1))   # HWC → CHW
    input_data = np.expand_dims(input_data, axis=0)     # 배치 차원 추가
 
    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: input_data})
    return outputs[0]

import torch
 
def export_to_onnx(model: torch.nn.Module, save_path: str) -> None:
    model.eval()
    dummy_input = torch.randn(1, 3, 224, 224)
 
    torch.onnx.export(
        model,
        dummy_input,
        save_path,
        export_params=True,
        opset_version=17,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}}  # 배치 크기 동적 처리
    )
    print(f"ONNX 모델 저장됨: {save_path}")

from typing import Any
import asyncio
from dataclasses import dataclass
from enum import Enum
 
class InferenceTarget(Enum):
    DEVICE = "device"
    CLOUD = "cloud"
 
@dataclass
class InferenceRequest:
    input_data: Any
    # 0.0~1.0: 입력 길이, 출력 엔트로피, 이전 실패율 등을 조합해 산출
    complexity_score: float
    requires_privacy: bool
    latency_budget_ms: float
 
class HybridInferenceRouter:
    COMPLEXITY_THRESHOLD = 0.7
    LATENCY_THRESHOLD_MS = 50.0
 
    def __init__(self, local_model: Any, cloud_client: Any) -> None:
        self.local_model = local_model
        self.cloud_client = cloud_client
 
    def route(self, request: InferenceRequest) -> InferenceTarget:
        # 프라이버시 요구 시 무조건 로컬
        if request.requires_privacy:
            return InferenceTarget.DEVICE
 
        # 레이턴시 예산이 타이트하면 로컬
        if request.latency_budget_ms < self.LATENCY_THRESHOLD_MS:
            return InferenceTarget.DEVICE
 
        # 복잡도가 높으면 클라우드
        if request.complexity_score > self.COMPLEXITY_THRESHOLD:
            return InferenceTarget.CLOUD
 
        return InferenceTarget.DEVICE
 
    async def infer(self, request: InferenceRequest) -> Any:
        target = self.route(request)
 
        if target == InferenceTarget.DEVICE:
            return await self.local_model.predict(request.input_data)
        else:
            return await self.cloud_client.predict(request.input_data)