1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
| """ YOLOv8 Xilinx Versal导出代码 ============================= """ import subprocess import os
def export_yolo_to_versal( model_name: str = "yolov8n", weights_path: str = "yolov8n.pt", output_dir: str = "./versal_export", quantize: bool = True, calibration_images: str = "./calibration_images", ): """ 将YOLOv8模型导出为Versal AI Edge兼容格式 步骤: 1. 导出为ONNX 2. AMD Quark量化 (INT8) 3. Vitis AI编译 Args: model_name: 模型名称 weights_path: PyTorch权重路径 output_dir: 输出目录 quantize: 是否量化 calibration_images: 量化校准图像目录 """ os.makedirs(output_dir, exist_ok=True) print(f"[1/3] 导出ONNX: {weights_path} → {output_dir}/model.onnx") cmd_onnx = [ "python3", "-c", f"from ultralytics import YOLO; " f"model = YOLO('{weights_path}'); " f"model.export(format='onnx', opset=13, simplify=True)" ] subprocess.run(cmd_onnx, check=True) onnx_path = weights_path.replace(".pt", ".onnx") if not quantize: print(f"导出完成(无量化): {onnx_path}") return onnx_path print(f"[2/3] AMD Quark INT8量化") cmd_quant = [ "quark", "quantize", "--input", onnx_path, "--output", f"{output_dir}/model_int8.onnx", "--calibration-data", calibration_images, "--quant-format", "int8", "--per-channel", ] subprocess.run(cmd_quant, check=True) print(f"[3/3] Vitis AI编译") cmd_compile = [ "vai_q_onnx", "--model", f"{output_dir}/model_int8.onnx", "--output_dir", output_dir, "--arch", "versal_ai_edge_gen2", ] subprocess.run(cmd_compile, check=True) compiled_model = f"{output_dir}/model_int8.xmodel" print(f"✅ 导出完成: {compiled_model}") print(f"\n=== 性能预估 (Versal AI Edge Gen2) ===") print(f" 模型: {model_name}") print(f" 量化: INT8") print(f" 预计FPS: 200-300 (取决于模型大小)") print(f" 预计延迟: 3-5ms") print(f" 预计功耗: 2-5W") return compiled_model
def export_ims_models(): """批量导出IMS所需YOLO模型""" ims_models = [ ("yolov8n-face", "yolov8n-face.pt", "人脸检测"), ("yolov8s-pose", "yolov8s-pose.pt", "人体姿态"), ("yolov8n-phone", "yolov8n-phone.pt", "手机检测"), ("yolov8s-seatbelt", "yolov8s-seatbelt.pt", "安全带检测"), ] for name, weights, desc in ims_models: print(f"\n{'='*50}") print(f"导出: {name} ({desc})") export_yolo_to_versal( model_name=name, weights_path=weights, output_dir=f"./versal_export/{name}", )
if __name__ == "__main__": export_ims_models()
|