LeViT-128S TensorRT FP16 Pipeline

facebook/levit-128S 모델을 TensorRT FP16 엔진으로 변환하고 고속 추론을 수행하는 파이프라인입니다.

환경

• GPU: RTX A6000 (Ampere)
• 정밀도: FP16
• 입력: 1x3x224x224 (고정)
• Python >= 3.10

디렉터리 구조

levit-128S-pipeline/
├── requirements.txt
├── configs/
│   └── preprocess.yaml
├── src/
│   ├── export_onnx.py
│   ├── build_trt_engine.py
│   ├── infer_trt.py
│   ├── benchmark_trt.py
│   └── validate_outputs.py
└── models/
    ├── levit_128s.onnx (생성됨)
    └── levit_128s_fp16_b1.engine (생성됨)

실행 순서

1. 환경 설정

python -m venv venv
source venv/bin/activate
uv pip install -r requirements.txt

2. ONNX 내보내기

timm에서 LeViT-128S 모델을 로드하고 ONNX 형식으로 내보냅니다.

python src/export_onnx.py

출력: models/levit_128s.onnx

3. TensorRT 엔진 빌드

ONNX 모델을 FP16 TensorRT 엔진으로 변환합니다.

python src/build_trt_engine.py

출력: models/levit_128s_fp16_b1.engine

4. 추론 실행

단일 이미지에 대해 추론을 수행합니다.

python src/infer_trt.py --image <이미지경로>

Top-5 클래스 인덱스와 확률이 출력됩니다.

5. 성능 벤치마크

워밍업과 반복 측정을 통해 지연시간과 FPS를 측정합니다.

python src/benchmark_trt.py

옵션:

• --num-warmup: 워밍업 반복 횟수 (기본: 50)
• --num-iters: 측정 반복 횟수 (기본: 500)

6. 출력 검증 (선택)

ONNXRuntime과 TensorRT 출력 간 코사인 유사도를 계산합니다.

python src/validate_outputs.py

스크립트 옵션

export_onnx.py

python src/export_onnx.py --out models/levit_128s.onnx --opset 17

build_trt_engine.py

python src/build_trt_engine.py \
    --onnx models/levit_128s.onnx \
    --engine models/levit_128s_fp16_b1.engine \
    --workspace_gb 8

infer_trt.py

python src/infer_trt.py \
    --engine models/levit_128s_fp16_b1.engine \
    --image <이미지경로> \
    --cfg configs/preprocess.yaml

benchmark_trt.py

python src/benchmark_trt.py \
    --engine models/levit_128s_fp16_b1.engine \
    --num-warmup 50 \
    --num-iters 500

validate_outputs.py

python src/validate_outputs.py \
    --onnx models/levit_128s.onnx \
    --engine models/levit_128s_fp16_b1.engine

전처리 설정

configs/preprocess.yaml에 정의된 ImageNet 표준 전처리를 사용합니다:

• 리사이즈: 224x224
• 정규화: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
• 중앙 크롭 적용

최적화 팁

• FP16 사용으로 RTX A6000에서 최대 성능 활용
• 고정 입력 형태(1x3x224x224)로 엔진 최적화
• 8GB 워크스페이스로 빌드 시간 단축 및 최적 tactics 탐색