PyTorch→ONNX→TFLite 변환 실패 8가지 해결

PyTorch에서 학습한 모델을 모바일/엣지로 가져가려면 보통 PyTorch → ONNX → (TF) → TFLite 파이프라인을 탑니다. 문제는 이 경로가 프레임워크 3개를 거치면서 연산자(op) 지원 범위, 텐서 레이아웃, 동적 shape, 양자화 규칙이 조금씩 달라 “어제 되던 모델이 오늘은 안 된다”가 흔하다는 점입니다.

이 글은 변환 실패를 에러 메시지 유형별로 8가지로 나눠, 재현 포인트와 해결책을 코드 중심으로 정리합니다. (예시는 분류/회귀 공통으로 적용되는 패턴 위주)

아래 단계는 최소한으로라도 먼저 고정해두면 디버깅이 훨씬 쉬워집니다.

• PyTorch export는 torch.onnx.export 또는 torch.onnx.dynamo_export 중 하나로 통일
• ONNX 검증은 onnx.checker + onnxruntime로 “ONNX 자체가 정상인지”부터 확인
• ONNX를 TF로 바꿀 때는 onnx-tf보다 최근엔 onnx2tf가 성공률이 높은 편
• TFLite 변환은 tf.lite.TFLiteConverter에서 supported_ops와 optimizations를 명시

참고로, 변환 과정에서 메모리 폭발이나 OOM이 같이 터질 때가 많은데, 그럴 땐 문제를 단순히 “변환 실패”로만 보지 말고 메모리 관점도 같이 보세요. 로컬 추론 OOM을 줄이는 패턴은 Transformers 로컬 LLM OOM 해결 - 4bit+KV 캐시 글의 접근(불필요 텐서/캐시 관리)도 도움이 됩니다.

변환 파이프라인 기준 코드(정상 기준선)

먼저 기준선을 하나 만들어두고, 이후 실패 케이스에서 어디가 깨졌는지 좁혀갑니다.

# 1) PyTorch -> ONNX
import torch
import onnx
import onnxruntime as ort

model = ...
model.eval()

dummy = torch.randn(1, 3, 224, 224)
onnx_path = "model.onnx"

torch.onnx.export(
    model,
    dummy,
    onnx_path,
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
)

onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)

# 2) ONNX Runtime로 ONNX 추론이 되는지 확인
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
out = sess.run(None, {"input": dummy.numpy()})
print("onnx ok", [o.shape for o in out])

ONNX가 여기서부터 실패하면 TFLite로 갈 이유가 없습니다. 반대로 ONNX가 정상인데 TF/TFLite에서 터지면 “연산자 지원/레이아웃/양자화” 쪽 문제일 확률이 큽니다.

실패 1) Unsupported operator 혹은 TFLite에서 SELECT_TF_OPS 요구

증상

• TF 변환은 되는데 TFLite 변환에서 Some ops are not supported by the native TFLite runtime 류 메시지
• 또는 변환 옵션에 SELECT_TF_OPS를 넣으라고 함

원인

TFLite 기본 런타임이 지원하는 op subset이 작아서, TF 그래프에 남아있는 op가 TFLite builtin으로 내려오지 못합니다. 특히 GridSample, NonMaxSuppression 변형, 일부 Resize/Gather 조합이 자주 문제를 만듭니다.

해결

1. 우선 “타협안”으로 SELECT_TF_OPS 허용
2. 그 다음 가능하면 모델을 수정해 builtin op로 떨어지게 리팩터링

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS,
]
# 일부 모델은 텐서리스트가 남아있어 이 옵션이 필요
converter._experimental_lower_tensor_list_ops = False

tflite_model = converter.convert()
open("model.tflite", "wb").write(tflite_model)

운영에서 SELECT_TF_OPS는 바이너리 크기 증가, 런타임 의존성 증가를 유발할 수 있어 최종 목표는 builtin-only입니다. 실패 op가 무엇인지 로그로 확인한 뒤, 해당 op를 다른 구현으로 바꾸는 것이 근본 해결입니다.

실패 2) Opset 불일치: opset_version 때문에 ONNX는 되는데 이후 변환이 깨짐

증상

• ONNX export는 성공
• ONNX 로딩/변환 단계에서 Unsupported opset 또는 특정 노드가 해석 불가

원인

opset_version이 너무 높거나(최신 op), 반대로 너무 낮아(구식 표현) 변환기에서 기대하는 형태와 달라집니다.

해결

• 보통 opset 13 또는 17이 타협점입니다.
• 변환기가 특정 opset에 더 안정적일 수 있으니, 실패 시 opset을 내려서 재시도합니다.

torch.onnx.export(
    model,
    dummy,
    "model_opset13.onnx",
    opset_version=13,
    input_names=["input"],
    output_names=["output"],
)

추가로, ONNX 그래프 최적화가 변환을 더 어렵게 만드는 경우도 있어 onnxsim 같은 simplifier 적용 여부를 케이스별로 비교하세요.

실패 3) 동적 shape 때문에 TF/TFLite에서 shape 추론 실패

증상

• 에러에 Unknown rank, None dimension, shape inference failed 류가 포함
• 배치나 시퀀스 길이를 동적으로 뒀더니 변환기에서 텐서 shape를 못 잡음

원인

ONNX는 동적 축을 꽤 유연하게 표현하지만, TF/TFLite 쪽은 변환 시점에 정적 shape를 더 선호합니다. 특히 TFLite는 정적 shape가 성공률이 높습니다.

해결

• 일단 배치/입력 길이를 고정한 “배포용 export”를 별도로 만듭니다.
• 또는 동적 축을 배치 하나만 남기고 나머지는 고정합니다.

# 배치만 동적, 나머지 고정
torch.onnx.export(
    model,
    dummy,
    "model_static_hw.onnx",
    opset_version=17,
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
)

모바일/엣지 배포라면 입력 크기를 고정하는 편이 성능/메모리 예측도 쉬워서 실무적으로도 유리합니다.

실패 4) 레이아웃 문제: NCHW와 NHWC 불일치로 결과가 틀리거나 변환이 실패

증상

• 변환은 되는데 결과가 완전히 다름
• 또는 Transpose가 과도하게 삽입되어 TFLite 변환이 실패/성능 급락

원인

PyTorch는 기본이 NCHW, TF/TFLite는 대부분 NHWC에 최적화되어 있습니다. 변환 과정에서 자동 transpose가 들어가며 그래프가 복잡해지고, 일부 경로에서 지원 op 조합이 깨집니다.

해결

• 가능하면 PyTorch 모델 입력부터 NHWC로 맞추는 전략(모델 구조에 따라 난이도 높음)
• 현실적으로는 “불필요 transpose를 줄이기” 위해 export 전후로 그래프를 단순화하거나, TF 변환 후 Transpose 패턴을 확인합니다.

PyTorch에서 입력을 바꾸는 단순 예시는 다음과 같습니다.

# 입력이 NHWC로 들어온다고 가정하고, 모델 앞단에서 NCHW로 바꿔 처리
class NHWCWrapper(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, x):
        # x: [N, H, W, C] -> [N, C, H, W]
        x = x.permute(0, 3, 1, 2).contiguous()
        return self.m(x)

wrapped = NHWCWrapper(model).eval()
dummy_nhwc = torch.randn(1, 224, 224, 3)

이 방식은 transpose를 “모델 바깥”이 아니라 “모델 안”으로 명시해 변환기의 예측 가능성을 높입니다.

실패 5) aten::/prim:: 노드가 ONNX에 남음(Export가 완전하지 않음)

증상

• ONNX 그래프에 aten:: 또는 prim:: 같은 PyTorch 전용 노드가 남아 변환기에서 거부
• 에러에 Exporting the operator ... to ONNX opset ... is not supported 류 포함

원인

해당 연산이 ONNX 표준 op로 매핑되지 않았거나, tracing 과정에서 그래프가 끊겼습니다. 커스텀 op, 조건 분기, 파이썬 제어 흐름이 섞이면 자주 발생합니다.

해결

• 가능하면 모델을 torch.nn 표준 op 조합으로 재작성
• 또는 torch.onnx.export 대신 torch.onnx.dynamo_export를 시도(모델에 따라 더 잘 잡음)
• 불가피하면 ONNX 커스텀 op를 정의해야 하는데, TFLite까지 가는 경로에서는 현실적으로 유지보수가 어렵습니다.

import torch

# PyTorch 2.x에서 시도 가능한 대안
exported = torch.onnx.dynamo_export(model, dummy)
exported.save("model_dynamo.onnx")

aten::가 남는 경우는 “변환 실패”라기보다 “export 설계 실패”에 가깝기 때문에, 이 단계에서 모델 구현을 정리하는 게 가장 빠른 해결입니다.

실패 6) 양자화(특히 INT8)에서 대표 데이터셋/입력 스펙 문제

증상

• FP32 TFLite는 되는데 INT8 변환에서 실패
• 에러에 representative_dataset 또는 quantization parameters 관련 문구

원인

완전 정수 양자화는 캘리브레이션용 대표 데이터가 필요하고, 입력 타입/범위가 요구사항을 만족해야 합니다. 또한 일부 op는 INT8 경로가 제한적이라 부분적으로만 양자화가 가능합니다.

해결

• 대표 데이터셋을 반드시 제공
• 입력/출력 타입을 명확히 지정
• 실패 시 “혼합 양자화(일부는 FP16/FP32)”로 타협

import tensorflow as tf
import numpy as np

def rep_data_gen():
    for _ in range(100):
        # 모델 입력 스펙에 맞게 생성
        x = np.random.rand(1, 224, 224, 3).astype(np.float32)
        yield [x]

converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = rep_data_gen

# 완전 INT8을 원할 때
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

tflite_int8 = converter.convert()
open("model_int8.tflite", "wb").write(tflite_int8)

INT8이 계속 실패하면, 우선 inference_input_type을 tf.float32로 두고 내부만 양자화되는지 확인한 뒤, 완전 INT8로 단계적으로 좁히는 것이 좋습니다.

실패 7) Resize/Interpolate 계열 변환 문제(align_corners 등)

증상

• 업샘플링이 들어간 모델에서 변환기 에러
• 또는 변환은 되는데 출력이 미세하게 달라져 품질이 크게 흔들림

원인

PyTorch interpolate의 옵션(align_corners, antialias)이 TF/TFLite의 ResizeBilinear/ResizeNearestNeighbor와 정확히 일치하지 않거나, ONNX 표현이 변환기에서 지원되지 않습니다.

해결

• 가능한 조합으로 제한: mode="nearest" 또는 표준 bilinear + align_corners=False
• export 전에 해당 레이어를 “지원되는 resize op”로 치환

import torch.nn.functional as F

class SafeResize(torch.nn.Module):
    def __init__(self, size):
        super().__init__()
        self.size = size

    def forward(self, x):
        # 변환기 호환성이 비교적 좋은 옵션으로 고정
        return F.interpolate(x, size=self.size, mode="bilinear", align_corners=False)

이 케이스는 실패도 실패지만, “성공했는데 품질이 깨지는” 형태가 더 위험합니다. 변환 후에는 반드시 샘플 입력 몇 개로 PyTorch vs TFLite 출력을 비교하세요.

실패 8) 그래프가 너무 커서 변환 중 메모리/시간 초과 또는 내부 크래시

증상

• 변환 프로세스가 죽거나, 매우 느려짐
• 빌드 환경에서 OOM, 혹은 변환 로그가 중간에 끊김

원인

대형 모델(특히 Transformer 계열)에서 상수 폴딩, 최적화, 캘리브레이션 단계가 메모리를 크게 먹습니다. 또한 변환 파이프라인이 중간 산출물을 크게 만들 수 있습니다.

해결

• 변환 단계별로 산출물을 저장하고(ONNX, SavedModel, TFLite) 단계 분리
• 불필요한 최적화 비활성화 또는 최소화
• FP16 변환으로 먼저 성공시킨 뒤 INT8로 확장

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")

# 먼저 FP16으로 성공률을 올려서 파이프라인을 고정
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

tflite_fp16 = converter.convert()
open("model_fp16.tflite", "wb").write(tflite_fp16)

대형 모델 변환은 “한 번에 끝내기”보다 “작게 쪼개서 어디서 폭발하는지”를 보는 게 핵심입니다. 이 접근은 네트워크 장애를 재현하고 원인을 줄여나가는 방식과 유사합니다. (예: AWS ALB 502·504 원인 - NLB·타임아웃·헬스체크에서 타임아웃/헬스체크를 분리해 진단하듯이, 변환도 단계별로 분리하면 원인 파악이 빨라집니다.)

변환 성공률을 올리는 체크리스트(실무용)

1) “ONNX가 정상인지”를 먼저 고정

• onnx.checker.check_model
• onnxruntime로 더미 입력 추론

2) 입력/출력 시그니처를 단순하게

• 입력은 가능하면 1개, dtype은 float32부터
• 동적 shape는 배치만 남기고 나머지는 고정

3) 변환 후 정확도/출력 비교를 자동화

• PyTorch 출력과 TFLite 출력의 max_abs_error, cosine similarity 등을 측정

import numpy as np

def max_abs_err(a, b):
    return float(np.max(np.abs(a - b)))

4) 실패 로그를 “op 이름”으로 분류

• Unsupported op면 SELECT_TF_OPS로 임시 통과 후 제거 작업
• shape면 정적화
• quant면 대표 데이터/입력 dtype부터 재점검

마무리

PyTorch → ONNX → TFLite 변환 실패는 대부분 “특정 연산자 1개” 또는 “shape/레이아웃/양자화 규칙 1개”에서 시작합니다. 중요한 건 한 번에 해결하려고 하기보다, ONNX 단계에서 정상 동작을 고정하고, 그 다음 TF/TFLite에서 깨지는 지점을 op 단위로 축소하는 것입니다.

위 8가지를 순서대로 점검하면, 흔한 변환 실패의 대부분은 재현 가능하고 수정 가능한 형태로 바뀝니다. 특히 Unsupported operator, 동적 shape, INT8 양자화 3가지는 실무에서 가장 많이 시간을 잡아먹는 포인트이니, 먼저 여기부터 체크해보세요.