Ohhnews

分类导航

$ cd ..
Baeldung原文

Triton Java API 入门指南：在 Java 中调用 NVIDIA 推理服务器

#triton#java#推理服务#目标检测#grpc

[LOADING...]

1. 简介

我们通过巧妙的工程实践，将复杂的机器学习模型部署到生产环境。尽管模型研究、训练和微调主要使用 Python 以及 PyTorch 或 TensorFlow 等框架完成，但企业后端通常使用 Java 构建。因此，我们需要一个健壮、可扩展的服务基础设施来弥合这一差距。

NVIDIA 的 Triton Inference Server 是一款开源模型服务软件，可标准化模型部署与执行。它提供统一的架构，能够托管几乎任何框架训练的模型，例如 ONNX、TensorFlow、PyTorch 以及 NVIDIA 高度优化的 TensorRT。Triton 为我们提供动态批处理、并发模型执行和内存管理，同时向客户端应用程序暴露简洁的 API。

在本教程中，我们将学习如何使用 Java 与 Triton Inference Server 交互，以演示图像中的目标检测。我们首先回顾可用的 API，然后使用 Docker 搭建本地 Triton 实例，最后编写一个简洁的 Java 客户端来执行预训练的 YOLO 模型。

2. Java API 概述

将 Triton 与 Java 应用程序集成主要有两种途径。

2.1. 进程内 Java API

进程内 Java API 使用 Java Native Interface (JNI) 绑定，直接与底层的 libtritonserver C 库通信。 这里，我们绕过网络协议，将 Triton 服务器实例直接嵌入 JVM 进程中。我们将这种方法用于低延迟边缘部署或网络资源匮乏的环境。

该 API 中的一些重要类包括：

  • TritonServer：表示嵌入式服务器实例。
  • TritonModel：表示已加载并可供推理的模型。
  • TritonRequest 和 TritonResponse：帮助我们在内存中来回传递张量。

2.2. Java 客户端 API（gRPC/HTTP）

对于大多数企业微服务架构，模型托管在集中式 CPU 或 GPU 集群上。在这种情况下，我们通过 gRPC 使用 Java 客户端 API 来访问模型。在底层，Triton 暴露了健壮的 protobuf 定义，Java 应用程序可以将其编译为强类型存根。

gRPC API 的关键类包括：

  • ManagedChannel：表示到 Triton 服务器的底层 gRPC 连接池。
  • InferenceServerBlockingStub：保存根据 Triton 的 protobuf 定义生成的同步客户端存根。
  • ModelInferRequest：封装我们的输入张量、形状和数据类型。
  • ModelInferResponse：给出包含计算预测结果的输出。
  • InferTensorContents：用于将原始基本类型（如浮点数或整数）安全地打包为字节流。

2.3. 官方仓库

对于更高级的用法，NVIDIA 提供了一个官方仓库，其中包含 client/src/java 下的实用封装。它提供了一个更简洁的 TritonClient 类，抽象掉样板式的 protobuf 生成，并支持异步推理和共享内存执行。 除此之外，Triton 的其他一些高级功能包括：

  • 异步推理：使用 gRPC 的异步存根处理高吞吐量、非阻塞请求。
  • 共享内存：允许 Triton 和 Java 客户端从同一系统内存（或 CUDA 内存）空间读写，消除 gRPC 的序列化和网络开销。
  • 字符串张量：将文本数据传递给 BERT 或 LLaMA 等 NLP 模型。

在本文中，我们将重点构建原生 gRPC Java 客户端 API，因为它是自定义企业环境中最常见的集成模式。

3. 使用 Docker 进行本地设置

我们将展示使用运行在 Triton Inference Server 上的 YOLO 预训练模型进行标准目标检测。为了在本地运行示例，我们将使用 Docker Desktop 运行一个活跃的 Triton Inference Server 实例。

3.1. 先决条件

以下是在基于 CPU 的本地机器上运行此设置所需的先决条件：

  • Docker Desktop。
  • Java 11 或更高版本。
  • 用于依赖管理的 Maven。
  • Python 3.10 及以上版本。

3.2. 完整项目结构

以下是该项目的完整目录结构：

$ shell
triton-java-yolo/
├── python/
│   ├── export_model.py              # It exports YOLO to ONNX 
│   └── requirements.txt             # Python dependencies (torch, ultralytics, onnx)
│
├── model_repository/                # It's a directory mounted into Docker Triton
│   └── yolo_onnx/
│       ├── 1/
│       │   └── model.onnx           # Compiled ONNX CPU runtime engine file
│       └── config.pbtxt             # Triton model configuration
│
├── src/
│   ├── main/
│   │   ├── proto/                   # Triton gRPC Protobuf definitions
│   │   │   ├── grpc_service.proto
│   │   │   └── model_config.proto
│   │   ├── java/
│   │   │   └── com/baeldung/triton/
│   │   │       ├── client/
│   │   │       │   └── TritonClientManager.java   
│   │   │       ├── yolo/
│   │   │       │   ├── ImagePreprocessor.java    
│   │   │       │   ├── YoloPostprocessor.java     
│   │   │       │   └── YoloInferenceRunner.java   
│   │   │       └── App.java                       
│   │   └── resources/
│   │       └── sample.jpeg           # Sample input image for object detection
│   └── test/
│       └── java/
│           └── com/baeldung/triton/
│               └── TritonInferenceLiveTest.java
│
├── pom.xml                          # Maven build file with gRPC & Protobuf plugins

3.3. 构建模型

首先，我们创建文件 export_model.py：

$ python
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
model.export(format="onnx", imgsz=640, dynamic=False)
print("Export complete: yolov8n.onnx generated.")

然后，我们构建 Python 虚拟环境，并使用它来运行文件 export_model.py：

$ shell
python export_model.py

这将下载预训练模型，然后在我们的工作目录中生成 ONNX 运行时 yolov8n.onnx：

$ shell
Downloading https://github.com/ultralytics/assets/releases/download/v8.4.0/yolovDownloading https://github.com/ultralytics/assets/releases/download/v8.4.0/yolov8n.pt to 'yolov8n.pt': 100% ━━━━━━━━━━━━ 6.2MB 42.1MB/s 0.1s
Ultralytics 8.4.34 🚀 Python-3.12.5 torch-2.2.2 CPU (Intel Core i9-9880H 2.30GHz)
YOLOv8n summary (fused): 72 layers, 3,151,904 parameters, 0 gradients, 8.7 GFLOPs
PyTorch: starting from 'yolov8n.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 84, 8400) (6.2 MB)
ONNX: starting export with onnx 1.22.0 opset 17...

3.4. 创建模型仓库

Triton 从称为模型仓库的严格格式化目录结构加载模型。因此，我们将创建 model_repository 目录结构，并移动和重命名 ONNX 文件：

$ shell
cp yolov8n.onnx ../model_repository/yolo_onnx/1/model.onnx

之后，我们将更新文件 config.pbtxt：

$ properties
name: "yolo_onnx"
backend: "onnxruntime"
max_batch_size: 0
input [
  {
    name: "images"
    data_type: TYPE_FP32
    dims: [ 1, 3, 640, 640 ]
  }
]
output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 1, 84, 8400 ]
  }
]

3.5. 运行服务器

准备好仓库后，让我们使用 Docker 容器启动 Triton Inference Server，并挂载我们的 model_repository 目录：

$ shell
docker run --platform linux/amd64 --rm \ -p 8000:8000 -p 8001:8001 -p 8002:8002 \ -v /absolute/path/to/model_repository:/models \ nvcr.io/nvidia/tritonserver:23.10-py3 \ tritonserver --model-repository=/models

这里，我们使用端口 8000 处理 HTTP REST 请求，端口 8001 用于 Java 应用程序中的 gRPC 端点，端口 8002 用于 Prometheus 指标。

我们可以通过查看日志来验证容器：

$ shell
=============================
== Triton Inference Server ==
=============================
NVIDIA Release 23.10 (build 72127154)
Triton Server Version 2.39.0
Copyright (c) 2018-2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
I0822 07:34:40.100142 1 server.cc:619] 
+-------------+-----------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Backend     | Path                                                            | Config                                                                                                                                                        |
+-------------+-----------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| onnxruntime | /opt/tritonserver/backends/onnxruntime/libtriton_onnxruntime.so | {"cmdline":{"auto-complete-config":"true","backend-directory":"/opt/tritonserver/backends","min-compute-capability":"6.000000","default-max-batch-size":"4"}} |
+-------------+-----------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
I0822 07:34:40.100196 1 server.cc:662] 
+-----------+---------+--------+
| Model     | Version | Status |
+-----------+---------+--------+
| yolo_onnx | 1       | READY  |
+-----------+---------+--------+

3.6. 配置 Protobuf

要运行 Java 客户端，我们需要设置 gRPC proto 和 Maven 依赖项。首先，我们将官方的 Triton proto 文件（来自 Triton Server 仓库的 grpc_service.proto 和 model_config.proto）放入 src/main/proto/：

$ shell
curl -L https://raw.githubusercontent.com/triton-inference-server/common/main/protobuf/grpc_service.proto -o src/main/proto/grpc_service.proto
curl -L https://raw.githubusercontent.com/triton-inference-server/common/main/protobuf/model_config.proto -o src/main/proto/model_config.proto
```## 4. Java 客户端

接下来，我们构建 Java 应用程序。

### 4.1. Maven 依赖项

首先，我们需要为该项目声明 Maven 依赖项。对于 gRPC，我们将使用核心库：[*grpc-netty-shaded*](https://mvnrepository.com/artifact/io.grpc/grpc-netty-shaded)、[*grpc-protobuf*](https://mvnrepository.com/artifact/io.grpc/grpc-protobuf) 和 [*grpc-stub*](https://mvnrepository.com/artifact/io.grpc/grpc-stub)，以及 [*javax.annotation-api*](https://mvnrepository.com/artifact/javax.annotation/javax.annotation-api)。为了自动将 .proto 文件编译为 Java 类，我们选择 [*os-maven-plugin*](https://mvnrepository.com/artifact/kr.motd.maven/os-maven-plugin) 和 [*protobuf-maven-plugin*](https://mvnrepository.com/artifact/org.xolstice.maven.plugins/protobuf-maven-plugin)。

以下是我们的依赖项：

```xml
<dependencies>
    <!-- gRPC & Protobuf Dependencies -->
    <dependency>
        <groupId>io.grpc</groupId>
        <artifactId>grpc-netty-shaded</artifactId>
        <version>${grpc.version}</version>
    </dependency>
    <dependency>
        <groupId>io.grpc</groupId>
        <artifactId>grpc-protobuf</artifactId>
        <version>${grpc.version}</version>
    </dependency>
    <dependency>
        <groupId>io.grpc</groupId>
        <artifactId>grpc-stub</artifactId>
        <version>${grpc.version}</version>
    </dependency>
    <dependency>
        <groupId>javax.annotation</groupId>
        <artifactId>javax.annotation-api</artifactId>
        <version>1.3.2</version>
    </dependency>
</dependencies>

接下来是我们的插件列表：

$ xml
<plugins>
    <!-- Compiles the .proto files in src/main/proto into Java classes -->
    <plugin>
        <groupId>org.xolstice.maven.plugins</groupId>
        <artifactId>protobuf-maven-plugin</artifactId>
        <version>0.6.1</version>
        <configuration>
            <protocArtifact>com.google.protobuf:protoc:${protobuf.version}:exe:${os.detected.classifier}</protocArtifact>
            <pluginId>grpc-java</pluginId>
            <pluginArtifact>io.grpc:protoc-gen-grpc-java:${grpc.version}:exe:${os.detected.classifier}</pluginArtifact>
        </configuration>
        <executions>
            <execution>
                <goals>
                    <goal>compile</goal>
                    <goal>compile-custom</goal>
                </goals>
            </execution>
        </executions>
    </plugin>
</plugins>

4.2. 建立连接

我们通过创建 ManagedChannel 并实例化一个阻塞存根（blocking stub）来发起通信：

$ java
ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 8001)
  .usePlaintext() 
  .build();
GRPCInferenceServiceBlockingStub blockingStub = GRPCInferenceServiceGrpc.newBlockingStub(channel);

4.3. 验证服务器健康状态

最佳实践是查询服务器的健康状态，以确保模型已成功加载，这样我们才能发送繁重的推理请求：

$ java
public boolean isServerLive() {
    try {
        ServerLiveRequest request = ServerLiveRequest.newBuilder().build();
        ServerLiveResponse response = blockingStub.serverLive(request);
        return response.getLive();
    } catch (Exception e) {
        return false;
    }
}

4.4. 在 Java 中准备输入张量

以下是使用标准 BufferedImage 的预处理逻辑：

$ java
InputStream is = ImagePreprocessor.class.getClassLoader().getResourceAsStream(resourcePath);
BufferedImage originalImage = ImageIO.read(is);
BufferedImage resizedImage = new BufferedImage(targetWidth, targetHeight, BufferedImage.TYPE_INT_RGB);
Graphics2D g = resizedImage.createGraphics();
g.drawImage(originalImage.getScaledInstance(targetWidth, targetHeight, Image.SCALE_SMOOTH), 0, 0, null);
g.dispose();

在将图像发送到 YOLOv8 模型之前，我们会对其进行预处理，以符合模型配置所要求的确切输入形状和格式。标准计算机视觉预处理包括调整大小、归一化和重新排列颜色通道。 YOLOv8 期望的输入张量形状为 [1, 3, 640, 640]。这对应于批大小为 1、3 个颜色通道（RGB）以及 640×640 的分辨率。 此外，数据必须采用平面格式（NCHW），即先存储所有红色像素，接着依次存储所有绿色和蓝色像素。最后，我们将像素提取并归一化为 float 数组 tensorData：

$ java
int totalPixels = targetWidth * targetHeight;
float[] tensorData = new float[3 * totalPixels];
int rOffset = 0, gOffset = totalPixels, bOffset = 2 * totalPixels;
for (int y = 0; y < targetHeight; y++) {
    for (int x = 0; x < targetWidth; x++) { 
        int rgb = resizedImage.getRGB(x, y); 
        int r = (rgb >> 16) & 0xFF;
        int gVal = (rgb >> 8) & 0xFF;
        int b = rgb & 0xFF;
        int index = y * targetWidth + x;
        tensorData[rOffset + index] = r / 255.0f;
        tensorData[gOffset + index] = gVal / 255.0f;
        tensorData[bOffset + index] = b / 255.0f;
    }
}

4.5. 执行推理

输入张量准备好后，我们现在构建 ModelInferRequest。

protobuf 定义表明，可以使用 outputTensor.getContents().getFp32ContentsList() 提取输出数据。Triton 通过对此列表留空来优化输出响应的性能，从而避免反序列化数十万个浮点数所带来的巨大性能开销。 相反，Triton 将底层的 C++ 内存缓冲区打包到 raw_output_contents 字段中，作为一个 ByteString。我们提取这些字节，并使用小端字节序（Little Endian）将其包装到 Java ByteBuffer 中，以便手动读取浮点数：

$ java
ModelInferRequest request = ModelInferRequest.newBuilder()
  .setModelName("yolo_onnx")
  .setModelVersion("1")
  .addInputs(inputTensor)
  .build();
ModelInferResponse response = blockingStub.modelInfer(request);
ByteString rawData = response.getRawOutputContents(0);
ByteBuffer buffer = rawData.asReadOnlyByteBuffer().order(ByteOrder.LITTLE_ENDIAN);
List resultList = new ArrayList<>(buffer.capacity() / 4);
while (buffer.hasRemaining()) {
    resultList.add(buffer.getFloat());
}

4.6. 非极大值抑制

YOLOv8 返回一个巨大的形状为 [1, C, N] 的矩阵，其中 C=84（COCO 数据集类别数），N>8000（图像中不同的锚框）。因此，这意味着当一个物体清晰可见时，多个重叠的锚框会为同一个物体报告高置信度检测结果。

为了防止应用程序为单个物体报告多个边界框，我们应用非极大值抑制（NMS）。NMS 会找出某个物体置信度最高的边界框，并抑制（丢弃）任何与其显著重叠的其他框。我们使用交并比（IoU）来确定重叠程度。如果一个置信度较低的框与置信度最高的框的重叠程度超过我们的阈值（例如 50%），则将其视为重复并移除：

$ java
if (current.classId == next.classId) {
    if (calculateIoU(current, next) > 0.5f) {
        suppressed[j] = true;
    }
}

4.7. 最终运行

现在，我们在主应用程序类中将所有组件整合在一起。首先导出 YOLO 模型。然后预处理示例猫图像。之后，通过 gRPC 运行推理，最后进行后处理并打印检测结果：

$ shell
mvn clean compile  
mvn exec:java -Dexec.mainClass="com.baeldung.triton.App"

它会产生干净、优化的输出，准确识别图像中的物体，且没有重复的边界框：

$ shell
Connecting to Triton Inference Server at localhost:8001
Triton Server is live and ready.
Building inference request for model: yolo_onnx
Sending inference request to Triton...
Received 705600 data points. Parsing bounding boxes...
Raw boxes found before NMS: 8
Detected [cat] (Confidence: 83.7%) at Box [xMin=2.0, yMin=53.2]
Final valid objects detected: 1

5. 测试

我们需要启动 Triton Inference Server Docker 容器，并确保它运行在 localhost:8001 上且已加载我们的 yolo_onnx 模型，以便运行实时推理测试。

在实时推理测试中，我们首先预处理示例猫图像。然后使用 YOLO 模型进行推理，该模型原生输出形状为 [1, 84, 8400] 的矩阵。接着，我们将其展平为 705,600 个元素，然后运行后处理器，使测试日志打印出经 NMS 处理后的猫检测结果：

$ java
public void givenValidImage_whenRunningInference_thenReturnsDetections() throws Exception {
    float[] inputTensor = ImagePreprocessor.preprocessFromResources("sample.jpeg", 640, 640);
    YoloInferenceRunner runner = new YoloInferenceRunner(clientManager.getStub(), "yolo_onnx");
    List<Float> outputs = runner.runInference(inputTensor);
    assertFalse(outputs.isEmpty(), "Inference output should not be empty");
    
    assertEquals(705600, outputs.size(), 
      "Output tensor should contain exactly 705,600 float elements");
    YoloPostprocessor.parseAndPrint(outputs);
}

6. 结论

在本文中，我们学习了 Triton Inference Server 及其与 Java 的集成。

简而言之，我们探索了 Triton Inference Server 如何在以 Python 为主导的机器学习生态系统与 Java 微服务之间提供强大的桥梁。通过使用 Docker，我们搭建了本地测试环境，构建了必要的基于 protobuf 的 gRPC 请求，并使用 ONNX CPU 运行时引擎成功执行了一次目标检测。

随着机器学习模型的规模和复杂度不断增长，使用像 Triton 这样的专用推理服务器可以确保我们的 Java 应用程序保持响应迅速、可扩展且易于维护。

一如既往，完整的代码示例可在 GitHub 上获取。

文章 Introduction to Triton Java API 最初发布于 Baeldung。