""" Contextra AI Analytics — RF-DETR TensorRT Export Script (v3.1) RF-DETR v1.9+ -> ONNX -> TensorRT Engine -> Triton Inference Server Triton Server 2.67.0 + TensorRT v10.16.00 호환 RF-DETR은 transformer 기반 모델로: - Batch 차원은 dynamic (1~8) - Detection output는 fixed (300 proposals, 91 classes) - opset 19+ 필요 (transformer layer 최적화) v3.1 수정 사항: - Nano max_proposals: 100 -> 300 (RFDETRNano도 300 proposals 출력) - config.pbtxt에서 모든 주석(#) 제거 - Triton protobuf 파서 호환성 확보 - Segmentation 모델 추가: RFDETRSegNano/Small/Medium/Large/XLarge/2XLarge - Detection 모델 추가: RFDETRBase - Pose는 별도 resolution(528) + keypoints 출력 (100x34x8) 사용 방법 (다른 시스템에 이전 시 한 번에 실행): cd /workspace/contextra-analytics/models python3 export_all_trt_v2.py Contextra AI Analytics — RF-DETR TensorRT Export Script (v3.2 - Fix divisibility) Failed models re-export with corrected resolutions. Fixes: - RFDETRBase: 512 -> 504 (divisible by 56=14*4, patch_size=14 windowed attn) - RFDETRSegSmall: 512 -> 504 (divisible by 24=12*2, patch_size=12 windowed attn) - RFDETRSegLarge: 704 -> 720 (divisible by 24) - RFDETRSegXLarge: 896 -> 888 (divisible by 24) - RFDETRSeg2XLarge: 1024 -> 1008 (divisible by 24) """ import os import subprocess import gc import logging from pathlib import Path logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logger = logging.getLogger("trt_multi_export") # ============================================================================ # Failed models with corrected resolutions # ============================================================================ FAILED_MODELS = { # Detection: patch_size=14, num_windows=4 -> divisible by 56 "RFDETRBase": { "type": "detection", "class": "RFDETRBase", "resolution": 504, # 504 / 56 = 9 (was 512) "max_proposals": 300, "num_classes": 91, }, # Segmentation: patch_size=12, num_windows=2 -> divisible by 24 "RFDETRSegSmall": { "type": "segmentation", "class": "RFDETRSegSmall", "resolution": 504, # 504 / 24 = 21 (was 512) "max_proposals": 300, "num_classes": 91, }, "RFDETRSegLarge": { "type": "segmentation", "class": "RFDETRSegLarge", "resolution": 720, # 720 / 24 = 30 (was 704) "max_proposals": 300, "num_classes": 91, }, "RFDETRSegXLarge": { "type": "segmentation", "class": "RFDETRSegXLarge", "resolution": 888, # 888 / 24 = 37 (was 896) "max_proposals": 300, "num_classes": 91, }, "RFDETRSeg2XLarge": { "type": "segmentation", "class": "RFDETRSeg2XLarge", "resolution": 1008, # 1008 / 24 = 42 (was 1024) "max_proposals": 300, "num_classes": 91, }, } DYNAMIC_BATCH = { "min_batch": 1, "opt_batch": 4, "max_batch": 8, } def _q(s): """Helper to produce a quoted string for pbtxt""" return '"' + s + '"' def generate_detection_config(model_name, input_size, max_proposals, num_classes): lines = [] lines.append("name: " + _q(model_name)) lines.append("platform: " + _q("tensorrt_plan")) lines.append("") lines.append("max_batch_size: 0") lines.append("") lines.append("input [") lines.append(" {") lines.append(" name: " + _q("input")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ 3, {}, {} ]".format(input_size, input_size)) lines.append(" }") lines.append("]") lines.append("") lines.append("output [") lines.append(" {") lines.append(" name: " + _q("dets")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ {}, 4 ]".format(max_proposals)) lines.append(" },") lines.append(" {") lines.append(" name: " + _q("labels")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ {}, {} ]".format(max_proposals, num_classes)) lines.append(" }") lines.append("]") lines.append("") return "\n".join(lines) def generate_segmentation_config(model_name, input_size, max_proposals, num_classes): lines = [] lines.append("name: " + _q(model_name)) lines.append("platform: " + _q("tensorrt_plan")) lines.append("") lines.append("max_batch_size: 0") lines.append("") lines.append("input [") lines.append(" {") lines.append(" name: " + _q("input")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ 3, {}, {} ]".format(input_size, input_size)) lines.append(" }") lines.append("]") lines.append("") lines.append("output [") lines.append(" {") lines.append(" name: " + _q("dets")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ {}, 4 ]".format(max_proposals)) lines.append(" },") lines.append(" {") lines.append(" name: " + _q("labels")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ {}, {} ]".format(max_proposals, num_classes)) lines.append(" },") lines.append(" {") lines.append(" name: " + _q("masks")) lines.append(" data_type: TYPE_FP32") lines.append(" dims: [ {}, {}, {} ]".format(max_proposals, input_size, input_size)) lines.append(" }") lines.append("]") lines.append("") return "\n".join(lines) def export_model(model_name, model_config): """Export a single failed model with corrected resolution""" mtype = model_config["type"] logger.info("=" * 60) logger.info("Re-exporting {} ({} model)".format(model_name, mtype)) logger.info("Resolution: {} -> {}".format( "original" if mtype == "detection" else "corrected", model_config["resolution"] )) base_dir = "/opt/contextra/models" triton_repo_dir = os.path.join(base_dir, "triton", model_name) version_dir = os.path.join(triton_repo_dir, "1") os.makedirs(version_dir, exist_ok=True) try: res = model_config["resolution"] logger.info("Loading {} (resolution={})...".format(model_name, res)) model_module = __import__("rfdetr", fromlist=[model_config["class"]]) ModelClass = getattr(model_module, model_config["class"]) model = ModelClass(resolution=model_config["resolution"]) logger.info("Model loaded: {}".format(type(model).__name__)) onnx_file_path = model.export( output_dir=base_dir, format="onnx", opset_version=19, dynamic_batch=True, shape=(model_config["resolution"], model_config["resolution"]), verbose=False ) onnx_path = str(onnx_file_path) logger.info("ONNX exported: {}".format(onnx_path)) final_engine_path = os.path.join(version_dir, "model.plan") logger.info("Converting to TensorRT Engine (FP16, dynamic batch)...") min_batch = DYNAMIC_BATCH["min_batch"] opt_batch = DYNAMIC_BATCH["opt_batch"] max_batch = DYNAMIC_BATCH["max_batch"] trtexec_cmd = [ "/usr/bin/trtexec", "--onnx={}".format(onnx_path), "--saveEngine={}".format(final_engine_path), "--fp16", "--verbose", "--minShapes=input:{}x3x{}x{}".format(min_batch, res, res), "--optShapes=input:{}x3x{}x{}".format(opt_batch, res, res), "--maxShapes=input:{}x3x{}x{}".format(max_batch, res, res), ] if mtype == "segmentation": trtexec_cmd.append("--memPoolSize=workspace:2048") else: trtexec_cmd.append("--memPoolSize=workspace:1024") logger.info("Running trtexec...") result = subprocess.run(trtexec_cmd, check=True, capture_output=True, text=True, timeout=900) logger.info("TensorRT engine created: {}".format(final_engine_path)) max_proposals = model_config["max_proposals"] num_classes = model_config["num_classes"] if mtype == "detection": config_content = generate_detection_config(model_name, res, max_proposals, num_classes) else: config_content = generate_segmentation_config(model_name, res, max_proposals, num_classes) config_path = os.path.join(triton_repo_dir, "config.pbtxt") with open(config_path, "w") as f: f.write(config_content) logger.info("Config written to: {}".format(config_path)) init_flag = os.path.join(base_dir, model_name, ".initialized") os.makedirs(os.path.dirname(init_flag), exist_ok=True) with open(init_flag, "w") as f: f.write("Exported at {} | TRT v10.16.00 | Triton 2.67.0\n".format(__import__('time').time())) logger.info("OK Successfully exported {}".format(model_name)) return True except Exception as e: logger.error("X Failed to export {}: {}".format(model_name, e), exc_info=True) return False def main(): """Re-export failed models with corrected resolutions""" logger.info("=" * 60) logger.info("RF-DETR TensorRT Export Pipeline (v3.2 - Retry)") logger.info("=" * 60) success_count = 0 total_count = len(FAILED_MODELS) for model_name, config in FAILED_MODELS.items(): if export_model(model_name, config): success_count += 1 gc.collect() logger.info("=" * 60) logger.info("Retry complete: {}/{} models".format(success_count, total_count)) logger.info("=" * 60) if __name__ == "__main__": main()