From c2dd067a56418c3739854b17da3b06a8f7b1d6e4 Mon Sep 17 00:00:00 2001
From: sulv <sulv6666@163.com>
Date: Thu, 10 Oct 2024 20:11:19 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=8A=A0=E8=BD=BD=E6=A8=A1?=
 =?UTF-8?q?=E5=9E=8B=E6=85=A2=EF=BC=8C=E9=A6=96=E6=AC=A1=E6=8E=A8=E7=90=86?=
 =?UTF-8?q?=E6=85=A2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main/java/com/ly/VideoInferenceApp.java   |  1 -
 .../com/ly/onnx/engine/InferenceEngine.java   | 28 ++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/ly/VideoInferenceApp.java b/src/main/java/com/ly/VideoInferenceApp.java
index 8de3422..b21b378 100644
--- a/src/main/java/com/ly/VideoInferenceApp.java
+++ b/src/main/java/com/ly/VideoInferenceApp.java
@@ -130,7 +130,6 @@ public class VideoInferenceApp extends JFrame {
             modelManager.loadModel(this);
             DefaultListModel<ModelInfo> modelList = modelManager.getModelList();
             ArrayList<ModelInfo> models = Collections.list(modelList.elements());
-
             for (ModelInfo modelInfo : models) {
                 if (modelInfo != null) {
                     boolean alreadyAdded = videoPlayer.getInferenceEngines().stream()
diff --git a/src/main/java/com/ly/onnx/engine/InferenceEngine.java b/src/main/java/com/ly/onnx/engine/InferenceEngine.java
index 0e62925..fbc51cb 100644
--- a/src/main/java/com/ly/onnx/engine/InferenceEngine.java
+++ b/src/main/java/com/ly/onnx/engine/InferenceEngine.java
@@ -12,6 +12,9 @@ import org.opencv.imgproc.Imgproc;
 
 import java.nio.FloatBuffer;
 import java.util.*;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 
 @Data
 public class InferenceEngine {
@@ -36,7 +39,30 @@ public class InferenceEngine {
     public InferenceEngine(String modelPath, List<String> labels) {
         this.modelPath = modelPath;
         this.labels = labels;
-        init();
+        initAsync();
+    }
+    // 异步执行模型初始化
+    public void initAsync() {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        executor.execute(()->{
+            init();
+            warmUp();
+        });
+    }
+
+    public void warmUp() {
+        // 提前执行一次空推理，用于初始化模型、CUDA上下文等
+        try {
+            float[] dummyInput = new float[(int) inputShape[0] * (int) inputShape[1] * (int) inputShape[2] * (int) inputShape[3]];
+            OnnxTensor inputTensor = OnnxTensor.createTensor(environment, FloatBuffer.wrap(dummyInput), inputShape);
+            String inputName = session.getInputInfo().keySet().iterator().next();
+            // 执行空推理
+            session.run(Collections.singletonMap(inputName, inputTensor));
+            inputTensor.close();
+            System.out.println("预热推理完成，首次推理性能已优化。");
+        } catch (Exception e) {
+            throw new RuntimeException("模型预热失败", e);
+        }
     }
 
     public void init() {