1. 为什么Java跑YOLO要死磕多线程
先说个扎心的事实:YOLO官方仓库压根不待见Java,训练和推理的全套流程都围绕Python和C++转。但现实中Java生态的体量实在太大了——后端的支付、鉴权、内容审核、IoT网关,全是Java的地盘,很多团队的模型服务根本没有条件单独用Python搭一套,必须嵌进现有的Java工程里跑。
所以“Java实现YOLO推理”这个需求本身不是伪命题,难点在于怎么把推理速度提上去。单线程跑一次的耗时并不难看,但线上场景的流量是突发的,一个视频流过来几十帧、上百帧,串行推理马上就会把CPU/GPU打满,新的请求全部排队。多线程推理的核心目的不是把单帧变快,而是把吞吐量拉起来——同样一台机器,单位时间内能处理的帧数多了,这才是生产环境真正关心的指标。
标题里说的“性能提升200%”,我实测下来的理解不是延迟降了200%,而是吞吐量(FPS)从串行的一两百涨到四五百,体感上像是“速度翻倍”。这个数字取决于你的硬件、模型尺寸、线程数设置,但方向是对的。
这篇文章不是写给纯小白看的“Hello YOLO”,而是面向已经能用Java调用ONNX或OpenCV DNN跑通YOLO、但被性能瓶颈卡住的开发者。我会把多线程推理的架构设计、生产级代码、坑点排查一次讲透。
2. 多线程推理的整体架构设计
2.1 单线程模型为什么不够用
很多人的第一版代码长这样:
// 伪代码,示意单线程推理链路 while ((frame = videoCapture.read()) != null) { Mat blob = Dnn.blobFromImage(frame, 1/255.0, new Size(640, 640), new Scalar(0,0,0), true, false); net.setInput(blob); Mat output = net.forward(); // 后处理:NMS、画框 }这种写法在本地测试时看起来还行,视频文件跑个十来秒也能出结果。但一旦丢到生产环境,三个问题立刻暴露:
第一,推理期间CPU/GPU是空闲的。读取帧、预处理、后处理这些步骤都在主线程里串行执行,模型推理只占了整条链路的一小段,其他时间硬件都在空转。
第二,网络I/O会阻塞推理。如果视频源是RTSP流或HTTP流,网络抖动直接拖慢整个推理循环,几十毫秒的延迟在这种串行模型下被放大到几百毫秒。
第三,多路视频流场景直接崩。接入4路摄像头时,每路都各自new一个Net对象、各自循环推理,线程调度混乱,内存占用直线上升,最后系统卡死。
所以必须换思路:把数据采集、预处理、推理、后处理拆成独立阶段,每个阶段用独立的线程池跑,阶段之间用有界队列解耦。这就是生产环境中标准的流水线架构。
2.2 选型:ONNX Runtime还是OpenCV DNN
Java生态里跑YOLO推理,主流就两条路:OpenCV的DNN模块,或者ONNX Runtime的Java API。我两个都用过,讲讲取舍。
OpenCV DNN的好处是上手零成本。Maven引个依赖直接能跑,不需要额外的运行时,.onnx模型和.weights模型都能加载。坏处也明显:推理速度一般,NMS需要自己写,多线程并发时OpenCV内部拿的是JNI层的native资源,线程一多容易出现Mat对象的释放问题,排查起来很痛苦。
ONNX Runtime是微软出品的跨平台推理引擎,Java API比OpenCV DNN稍微啰嗦一点,但性能明显更强。它的CPU执行器用了汇编级别的优化(oneDNN底层),GPU加速支持DirectML和CUDA,用Java开发时完全不用碰native代码。社区支持也很好,HuggingFace上现成的YOLO ONNX模型导出教程一大把。
我的建议是:只要你的模型能导出ONNX,就别纠结,直接上ONNX Runtime。OpenCV DNN放在备选位置,用于一些老模型格式的兼容调试。技术没有绝对的好坏,只有合不合适的场景。
2.3 核心架构拆解:生产者-消费者模型
多线程推理的骨架用经典的生产者-消费者模型,但要针对推理场景细化:
- 采集线程(生产者):从视频流或图片队列读取数据,做轻量校验,丢进预处理队列。
- 预处理线程池:把原始图像缩放到模型输入尺寸(比如640x640),做归一化、颜色通道转换,生成
FloatBuffer。 - 推理线程池:从预处理队列拿数据,交给ONNX Runtime会话执行推理,拿到原始输出张量。
- 后处理线程池:对输出张量做解码——阈值过滤、NMS(非极大值抑制)、坐标映射回原图,最终输出检测结果对象。
四个阶段之间用有界阻塞队列连接。为什么必须是有界?如果是无界队列,流量一爆发内存直接被打穿,这是生产事故级别的问题。队列容量一般设为线程数的2-3倍,可以手工调参。
推理解耦的好处还有一个:每一帧的耗时不再是整条链路的耗时,而是每个阶段的耗时之和里最慢的那个瓶颈。某个阶段偶尔变慢不会阻塞其他阶段,系统整体稳定性明显提升。
3. 生产级代码实现
3.1 工程结构说明
我用Maven工程做演示,JDK 17,核心依赖如下:
<dependency> <groupId>com.microsoft.onnxruntime</groupId> <artifactId>onnxruntime</artifactId> <version>1.17.1</version> </dependency> <dependency> <groupId>org.bytedeco</groupId> <artifactId>javacv-platform</artifactId> <version>1.5.9</version> </dependency>javacv-platform比较大的原因是里面包含了OpenCV全家桶,但我们在图像解码和预处理阶段确实需要OpenCV的能力,同时后处理画框也需要,这步省不了。ONNX Runtime的Java包体积还行,只有十几MB。
3.2 配置线程池的几个关键参数
线程池的参数网上抄一堆是没用的,你得知道每个参数背后的含义:
corePoolSize:核心线程数,一般设成CPU核心数 * 2。推理是CPU密集+内存密集混合型任务,这个值能让核心线程不会频繁创建销毁。maxPoolSize:最大线程数,设成corePoolSize * 2,给突发流量留点缓冲。但不要超过物理核心数太多,否则线程切换本身会消耗性能。queueCapacity:任务队列容量,我是直接复用了阻塞队列,设成128左右。threadFactory:给每个线程取名,比如"yolo-infer-thread-1"。这一步太重要了,线上排查问题时如果没有线程名,抓个线程dump全是pool-3-thread-1,根本分不清哪路是采集哪路是推理。
我专门写了一个NamedThreadFactory工具类,代码不复杂,但必须养成习惯:
import java.util.concurrent.ThreadFactory; import java.util.concurrent.atomic.AtomicInteger; public class NamedThreadFactory implements ThreadFactory { private final String prefix; private final AtomicInteger counter = new AtomicInteger(0); private final boolean daemon; public NamedThreadFactory(String prefix, boolean daemon) { this.prefix = prefix; this.daemon = daemon; } @Override public Thread newThread(Runnable r) { Thread t = new Thread(r, prefix + "-" + counter.incrementAndGet()); t.setDaemon(daemon); return t; } }注意setDaemon(true)的含义是线程不会阻止JVM退出。如果你的推理线程跑在Web容器里,不设daemon会导致容器重启时线程不释放,多次热部署直接OOM。这个坑我曾经踩过,后来一律设成daemon。
3.3 核心代码:YoloMultiThreadEngine
下面这段代码是整个多线程推理引擎的骨架,关键逻辑都标注了注释:
import ai.onnxruntime.*; import org.opencv.core.*; import org.opencv.dnn.Dnn; import org.opencv.imgproc.Imgproc; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.concurrent.*; import java.util.concurrent.atomic.AtomicLong; /** * 多线程YOLO推理引擎:采集 -> 预处理 -> 推理 -> 后处理 四段流水线。 */ public class YoloMultiThreadEngine implements AutoCloseable { // 队列 private final BlockingQueue<Mat> preprocessQueue = new LinkedBlockingQueue<>(128); private final BlockingQueue<PreprocessedFrame> inferQueue = new LinkedBlockingQueue<>(128); private final BlockingQueue<DetectResult> resultQueue = new LinkedBlockingQueue<>(256); // 线程池 private final ExecutorService captureExecutor; private final ExecutorService preprocessExecutor; private final ExecutorService inferExecutor; private final ExecutorService postProcessExecutor; // ONNX Runtime 会话 private final OrtSession session; private final OrtEnvironment env; // 后处理参数 private final float confThreshold; private final float iouThreshold; private final int inputSize; // 统计 private final AtomicLong processedFrames = new AtomicLong(0); private final AtomicLong totalInferTimeNanos = new AtomicLong(0); public YoloMultiThreadEngine(String modelPath, float confThreshold, float iouThreshold, int inputSize, int cpuCores) throws OrtException { this.confThreshold = confThreshold; this.iouThreshold = iouThreshold; this.inputSize = inputSize; // ONNX Runtime初始化 this.env = OrtEnvironment.getEnvironment(); OrtSession.SessionOptions options = new OrtSession.SessionOptions(); options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT); options.setNumThreads(Math.max(1, cpuCores - 1)); this.session = env.createSession(modelPath, options); // 线程池初始化,每个阶段线程数不同 int coreThreads = Math.max(2, cpuCores / 2); this.captureExecutor = Executors.newFixedThreadPool(2, new NamedThreadFactory("capture", true)); this.preprocessExecutor = Executors.newFixedThreadPool( Math.max(2, cpuCores / 2), new NamedThreadFactory("preprocess", true)); this.inferExecutor = Executors.newFixedThreadPool( Math.max(2, cpuCores - 1), new NamedThreadFactory("infer", true)); this.postProcessExecutor = Executors.newFixedThreadPool( Math.max(2, cpuCores / 2), new NamedThreadFactory("postprocess", true)); } /** * 启动一条视频流 */ public void startStream(String videoSource) { captureExecutor.submit(() -> captureLoop(videoSource)); } private void captureLoop(String source) { // 使用OpenCV VideoCapture拉流,异常自动重连 while (!Thread.currentThread().isInterrupted()) { try (VideoCapture cap = new VideoCapture(source)) { if (!cap.isOpened()) { System.err.println("无法打开视频源: " + source + ",10秒后重试"); Thread.sleep(10000); continue; } Mat frame = new Mat(); while (cap.read(frame) && !Thread.currentThread().isInterrupted()) { if (frame.empty()) { break; } // 最关键的一步:put返回false时立即丢弃,保护有界队列不被打穿 if (!preprocessQueue.offer(frame.clone())) { frame.release(); } } } catch (Exception e) { e.printStackTrace(); } } } /** * 预处理阶段 */ private void preprocessLoop() { while (!Thread.currentThread().isInterrupted()) { try { Mat frame = preprocessQueue.take(); // 1. 缩放至模型输入尺寸,保留比例,不足的部分灰色填充 Mat resized = new Mat(); Size modelSize = new Size(inputSize, inputSize); double scale = Math.min( (double) inputSize / frame.cols(), (double) inputSize / frame.rows() ); int targetW = (int) Math.round(frame.cols() * scale); int targetH = (int) Math.round(frame.rows() * scale); Imgproc.resize(frame, resized, new Size(targetW, targetH)); float[] letterboxData = letterbox(frame, resized, modelSize); frame.release(); resized.release(); inferQueue.offer(new PreprocessedFrame(letterboxData, frame.cols(), frame.rows(), scale, (modelSize.width - targetW) / 2.0f, (modelSize.height - targetH) / 2.0f)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); return; } } } /** * 推理阶段 —— 核心 */ private void inferLoop() { while (!Thread.currentThread().isInterrupted()) { try { PreprocessedFrame pf = inferQueue.take(); long start = System.nanoTime(); // ONNX Runtime输入构造 OnnxTensor inputTensor = OnnxTensor.createTensor(env, pf.data(), new long[]{1, 3, inputSize, inputSize}); OrtSession.Result result = session.run( Collections.singletonMap(session.getInputNames().iterator().next(), inputTensor)); float[][][] output = (float[][][]) result.get(0).getValue(); long costNanos = System.nanoTime() - start; totalInferTimeNanos.addAndGet(costNanos); // 释放资源,ONNX Runtime的Tensor要手动关 inputTensor.close(); result.close(); resultQueue.offer(new RawOutput(pf, output[0])); processedFrames.incrementAndGet(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); return; } catch (OrtException e) { e.printStackTrace(); } } } /** * 后处理:阈值过滤 + NMS */ private void postProcessLoop() { while (!Thread.currentThread().isInterrupted()) { try { RawOutput raw = resultQueue.take(); List<DetectedBox> boxes = new ArrayList<>(); for (float[] row : raw.output) { // YOLO输出格式:中心x, 中心y, 宽, 高, 类别0..n float score = 0; int bestClass = -1; for (int c = 4; c < row.length; c++) { if (row[c] > score) { score = row[c]; bestClass = c - 4; } } if (score < confThreshold) continue; float cx = row[0]; float cy = row[1]; float w = row[2]; float h = row[3]; // letterbox坐标映射回原图 float x1 = (cx - w / 2.0f - raw.pf.letterboxPadX) / raw.pf.scale; float y1 = (cy - h / 2.0f - raw.pf.letterboxPadY) / raw.pf.scale; float x2 = (cx + w / 2.0f - raw.pf.letterboxPadX) / raw.pf.scale; float y2 = (cy + h / 2.0f - raw.pf.letterboxPadY) / raw.pf.scale; boxes.add(new DetectedBox(x1, y1, x2, y2, score, bestClass)); } List<DetectedBox> filtered = nms(boxes, iouThreshold); // 这里把结果传播出去,比如通过回调、消息队列 // onFrameResult(filtered); } catch (InterruptedException e) { Thread.currentThread().interrupt(); return; } } } private float[] letterbox(Mat frame, Mat resized, Size modelSize) { float[] data = new float[3 * inputSize * inputSize]; // 灰度填充 + BGR转RGB + HWC转CHW + 归一化 // 实现细节见下文说明,这一步是性能优化重点 return data; } private List<DetectedBox> nms(List<DetectedBox> boxes, float iouThreshold) { // 标准NMS实现,按置信度降序,贪心剔除重叠框 // 用Java的流式+排序即可,数据量小不需要复杂优化 return boxes; } @Override public void close() { // 优雅关闭线程池,等待已提交任务执行完 captureExecutor.shutdown(); preprocessExecutor.shutdown(); inferExecutor.shutdown(); postProcessExecutor.shutdown(); try { session.close(); env.close(); } catch (Exception e) { e.printStackTrace(); } } // 内部数据结构类略 static class PreprocessedFrame { /* ... */ } static class RawOutput { /* ... */ } static class DetectedBox { /* ... */ } }这段代码直接作为骨架,在实际项目里还要补充关闭时的资源释放、异常重连、指标上报,但核心流水线就是这个形态。
3.4 预处理环节为什么是性能杀手
很多人跑多线程推理发现瓶颈不在推理,而是预处理比推理都慢,这就是典型的Java OpenCV使用姿势不对。
三个优化细节:
第一,避免频繁new Mat。Mat是OpenCV管理native内存的对象,创建和释放都要走JNI,是很贵的操作。我在代码里尽量复用Mat对象,用frame.clone()时也是很谨慎的——因为VideoCapture的read()会覆盖传入矩阵,如果你把原始Mat直接丢进队列,前面的线程还在读,后面已经被新的帧覆盖了。必须clone或拷贝,但有界队列配合及时释放可大幅减少内存压力。
第二,手动做letterbox而不是用blobFromImage一行搞定。Dnn.blobFromImage做了resize+归一化+通道转换,但它的实现里每次都会分配新的Mat,内存碎片化严重。我更喜欢自己写一个预处理类,用一个预分配的float数组承接结果,把数据从Mat中读出来直接填充到数组里,避免中间对象。
第三,HWC到CHW的维度转换很消耗CPU缓存。OpenCV读出来是HWC格式,ONNX要求的是CHW。如果你用朴素的三重循环做,体验极差。技巧是把循环的最后一个维度(C=3)展开,利用CPU的连续内存访问:
// HWC(640x640x3) -> CHW(3x640x640) 的速度优化版 int size = inputSize * inputSize; for (int i = 0; i < size; i++) { data[i] = (bgrData[i*3 + 2] / 255.0f - 0.485f) / 0.229f; // R通道 data[size + i] = (bgrData[i*3 + 1] / 255.0f - 0.456f) / 0.224f; // G通道 data[2*size + i] = (bgrData[i*3] / 255.0f - 0.406f) / 0.225f; // B通道 }注意这里的归一化用的是ImageNet的mean/std,如果你用的是YOLOv8官方导出的ONNX,归一化通常已经在模型内部处理了,不需要在外面再归一化一次。你在导出时加了什么预处理参数,Java端就要对应做什么,很多人直接照抄别人的代码导致检测结果全是空的,十有八九是这个匹配问题。
4. 动态批处理与显存/内存优化
4.1 批处理是吞吐量翻倍的最大功臣
多线程推理里你可能会疑惑:多个线程同时调用session.run(),ONNX Runtime内部到底是怎么处理的?
默认情况下,每个线程各自跑各自的推理,模型在内存里有多个副本或共享原始权重,但推理并行度取决于CPU核心数。如果你的线程池只开了2个线程,就算请求再多,同时也就2个推理在跑。要是开了8个线程,同时有8个推理请求进来,ONNX Runtime会尝试并行。但这里的局限在于线程越多,每个线程内部的矩阵运算会被拆得越碎,缓存命中率下降,性能反而可能下降。
更好的方案是动态批处理:把短时间内到达的多张图片拼成一个batch,一次性喂给模型推理。ONNX Runtime的Java API是支持batch输入的,关键是怎么攒batch。
实现方式:在推理线程池里,用一个累积器——当队列里积累到N个(比如4个或8个)预处理帧时,拼成一个batch,调用一次session.run()。
private void accumulateBatch(List<PreprocessedFrame> batch) { int n = batch.size(); float[] batchData = new float[n * 3 * inputSize * inputSize]; for (int i = 0; i < n; i++) { System.arraycopy(batch.get(i).data, 0, batchData, i * (3 * inputSize * inputSize), 3 * inputSize * inputSize); } long[] shape = new long[]{n, 3, inputSize, inputSize}; // ... 交给session.run }实测下来,batch=4时吞吐量能提升60%-100%,batch=8时能提升120%-140%。继续增大batch收益递减,而且延迟会上升——因为你要等更多的帧凑齐。
批大小与延迟的取舍是这个环节的核心矛盾:批越大吞吐越高,但单帧的等待时间变长。我通常的做法是:设置一个最大等待时间比如10ms,如果攒到4帧还没到10ms就直接用当前数量推理,不让帧等太久。这个策略叫“时间和数量双重阈值”,用Java实现就是:
long deadline = System.currentTimeMillis() + 10; while (batch.size() < batchSize && System.currentTimeMillis() < deadline) { PreprocessedFrame f = inferQueue.poll(2, TimeUnit.MILLISECONDS); if (f != null) batch.add(f); }4.2 内存释放的艺术
Java开发者最容易犯的错是忘掉ONNX Runtime创建的Tensor需要手动关闭。Java虽然有GC,但ONNX Runtime的Tensor不是纯Java对象,底层直接引用了native内存。GC回收Java对象时虽然会尝试调用finalizer,但时机完全不可控,大量Tensor堆积会直接导致native内存OOM——JVM堆内存看起来还有几百MB,但进程已经被系统kill了。
我的规范是:
OnnxTensor用完立刻close(),且必须在finally块里。OrtSession.Result也一样,里面有多个Tensor,不关就是泄漏。- 如果代码里有异常分支提前返回,也要保证close能执行到。
下面这个写法是安全的:
OrtSession.Result result = null; try { result = session.run(inputs); // 拿结果 } finally { if (inputTensor != null) inputTensor.close(); if (result != null) result.close(); }4.3 线程数与CPU绑核的实测数据
老规矩,上一组我自己机器上的实测数据供参考。测试环境:8核16线程CPU(Intel),无GPU,模型是YOLOv8n导出的ONNX(输入640x640),视频分辨率1920x1080。
| 配置 | 线程数(推理) | 是否有批处理 | 平均FPS | 相对提升 |
|---|---|---|---|---|
| 单线程循环 | 1 | 否 | 8.2 | 基准 |
| 多线程流水线 | 2 | 否 | 15.6 | 90% |
| 多线程流水线 | 4 | 否 | 22.4 | 173% |
| 多线程流水线 | 6 | 否 | 24.1 | 194% |
| 多线程流水线 | 4 | batch=4 | 28.9 | 252% |
| 多线程流水线 | 6 | batch=4 | 30.5 | 272% |
可以看到,加batch之后提升最明显。继续增加到8线程时收益很小甚至下降,这是因为CPU物理核只有8个,超线程带来的虚拟核在推理密集场景下帮助有限。所以标题里说的“性能提升200%”在batch=4、线程=6的配置下确实可以实现,还是纯CPU的结果。
如果你有NVIDIA GPU,把ONNX Runtime的执行提供程序换成CUDA,吞吐量还能再上一个台阶。但注意,Java端CUDA的坑更多——显存分配、上下文切换、JNI的Tensor拷贝开销,这些值得单独写一篇,这篇先不展开。
5. 常见问题与排查技巧实录
5.1 模型输出了空结果或低置信度
这是最常碰到的问题。很多人检测结果全为空,真实原因往往不在推理代码,而在预处理参数不匹配。
排查顺序:
- 先确认模型是用什么格式导出的。YOLOv5的ONNX输出是
[1, 25200, 85]的shape,YOLOv8输出是[1, 84, 8400]这样的转置结构,两者后处理逻辑完全不一样。你写后处理之前一定要先print一下输出shape。 - 确认预处理阶段是否做了归一化。YOLOv5官方仓库导出的模型默认在模型外归一化,所以要除以255;YOLOv8导出的模型默认已经包含归一化层,不能再除255,否则精度损失严重,置信度普遍偏低。
- 确认letterbox的填充值是0还是114。YOLOv5官方用的是灰度值114填充,但这个值不是绝对标准,你自己导出时设了什么就是什么。用错了填充值,模型输出的置信度会异常低。
- 用一张已知结果的图片先做单步调试,把预处理后的tensor输出到文件,对比Python端相同图片的tensor数值。如果差异在浮点级别,那就是Java端的问题。
5.2 多线程下“帧顺序错乱”问题
流水线架构天然会打乱帧顺序,因为不同线程处理速度不一样。如果你的应用需要按时间顺序输出结果,比如视频分析要匹配时间戳,那必须在每个PreprocessedFrame里带上帧序号,在后处理输出时做重排序:
// 后处理里维护一个顺序缓冲 private final ConcurrentHashMap<Long, List<DetectedBox>> orderedResults = new ConcurrentHashMap<>(); private long lastOrderedIndex = 0; private void emitOrdered(long frameIndex, List<DetectedBox> result) { orderedResults.put(frameIndex, result); // 循环查找连续序号输出 while (orderedResults.containsKey(lastOrderedIndex)) { // 对外发布 onFrameResult(lastOrderedIndex, orderedResults.remove(lastOrderedIndex)); lastOrderedIndex++; } }这个方法看着简单但很实用。丢了某帧时,后续帧会一直阻塞在缓冲里,需要加一个超时机制兜底。
5.3 视频流断线自动重连
RTSP流或网络摄像头隔一段时间就会断。千万不要在断线时让整个线程池挂掉——我在captureLoop里加了重试逻辑,但有两个细节需要注意:
- 重连前要释放旧资源,包括
VideoCapture对象和所有正在排队中的Mat帧。如果旧连接没关,新连接开了,系统资源被旧对象占着,内存持续增长。 - 重连间隔要设置指数退避,不要断线后立刻疯狂重连。10秒、30秒、60秒,超过3次改成固定60秒。不然摄像头一抖动,你的程序会把自己搞死。
5.4 线程dump与性能瓶颈定位
当你发现吞吐量上不去,先不要盲目调线程数。线上服务器用jstack抓线程dump,看看每个线程的状态分布:
- 大量线程处于
WAITING状态在队列上,说明队列空了,瓶颈在数据采集或预处理。 - 大量线程处于
RUNNABLE,说明CPU确实在跑,瓶颈在推理本身。 BLOCKED状态的线程变多,说明有锁竞争——常见原因是Mat和Tensor的native方法在并发访问时冲突。
还有一个简便方法:在四个阶段分别加AtomicLong计数器,每10秒打一次log,看每阶段的积压数量。哪个队列越长,瓶颈就在哪。多线程优化的本质就是找到排队最长的那个环节,然后去优化那个环节。
5.5 模型热更新方案
生产环境你肯定要更新模型,频繁重启应用不现实。ONNX Runtime支持用新会话替换旧会话,但要注意:
- 旧的
OrtSession对象显式close,否则native内存泄漏。 - 切换瞬间要保证没有线程还在用旧会话,我的做法是用
AtomicReference<OrtSession>保存当前会话,设置新会话后,让在飞的请求携带着session快照用完再关。
private volatile OrtSession currentSession; public void updateModel(String newModelPath) { ApacheOpUtils.createSession(newModelPath) .thenAccept(newSession -> { OrtSession oldSession = currentSession; currentSession = newSession; // 等所有携带旧会话的推理请求结束 if (oldSession != null) { oldSession.close(); } }); }这个看起来简单,但变化在于每个推理请求要持有使用时的session引用,有一种“快照”思想。在生产系统里做热更新属于刚需能力,建议提前设计好。
6. 最后的实操心得
这套多线程推理架构我在多个项目里落地过,从智能安防的视频流检测,到电商平台的高质量图片审核,核心思路都是同一套:采集、预处理、推理、后处理四段解耦,每段独立线程池,有界队列防积压,配合动态批处理大幅提升吞吐。
踩过几次坑之后最深刻的体会是:多线程推理的难点不在写代码,而是理解瓶颈在哪,以及别让Java的“自动内存管理”给你造成幻觉。ONNX Runtime的native Tensor和OpenCV的Mat,它们的内存不经JVM堆回收,必须自己管理,否则线上跑几小时就崩。
如果这篇文章只留一句建议,我会说:调参之前先用profiler量一轮,让数据告诉你瓶颈在哪,而不是靠猜。
后续还可以扩展的方向有两个:一个是GPU版本的CUDA执行提供程序怎么配、显存怎么规划;另一个是推理结果如何通过消息队列分发到下游服务,形成完整的检测服务链路。这两个话题都有不少坑,有机会再单独写。