如何用JAVA封装跨框架推理接口以适配不同大模型后端运行环境

JAVA 封装跨框架推理接口，适配不同大模型后端运行环境

大家好，今天我们来聊聊如何使用 JAVA 封装跨框架推理接口，以适配不同的大模型后端运行环境。随着大模型技术的飞速发展，涌现出了各种不同的推理框架，例如 TensorFlow Serving, Triton Inference Server, ONNX Runtime 等。在实际应用中，我们可能需要根据不同的需求和场景选择不同的推理后端。为了避免代码的重复编写和维护，我们需要一个统一的接口来访问这些不同的后端。

1. 问题分析与设计目标

在构建跨框架推理接口之前，我们需要明确需要解决的问题和设计目标。

问题：

• 框架差异性： 不同的推理框架具有不同的 API 和数据格式，直接使用会增加代码的复杂性和维护成本。
• 环境依赖性： 某些框架可能依赖特定的硬件或软件环境，导致部署困难。
• 代码冗余： 为每个框架编写单独的推理代码会导致大量冗余，不利于代码复用和维护。

设计目标：

• 统一接口： 提供一个统一的 JAVA 接口，屏蔽底层框架的差异。
• 可扩展性： 易于添加新的推理框架支持。
• 灵活性： 允许用户配置不同的后端实现。
• 高性能： 尽量减少封装带来的性能损耗。

2. 总体架构设计

为了实现上述目标，我们可以采用以下架构：

+---------------------+     +---------------------+     +---------------------+
|  Client Application | --> |   Inference Service  | --> | Backend Implementation |
+---------------------+     +---------------------+     +---------------------+
                                   |
                                   | (Interface Definition)
                                   |
     +---------------------+     +---------------------+     +---------------------+
     | TensorFlow Backend |     |   Triton Backend    |     |  ONNX Runtime Backend |
     +---------------------+     +---------------------+     +---------------------+

• Client Application： 客户端应用程序，通过统一的 Inference Service 接口进行推理。
• Inference Service： 定义推理服务的通用接口，负责接收请求、调用后端实现并返回结果。
• Backend Implementation： 不同的后端实现，例如 TensorFlow Backend, Triton Backend, ONNX Runtime Backend 等，负责与具体的推理框架进行交互。

3. 代码实现

接下来，我们通过代码示例来演示如何实现这个架构。

3.1 定义通用接口 InferenceService

首先，我们定义一个通用的 InferenceService 接口，该接口包含 infer 方法，用于接收输入数据并返回推理结果。

import java.util.Map;

public interface InferenceService {
    /**
     * 执行推理.
     *
     * @param input 输入数据，使用 Map 存储，key 为输入名称，value 为输入数据.
     * @return 推理结果，使用 Map 存储，key 为输出名称，value 为输出数据.
     * @throws InferenceException 如果推理过程中发生错误.
     */
    Map<String, Object> infer(Map<String, Object> input) throws InferenceException;
}

3.2 定义异常类 InferenceException

为了统一处理推理过程中可能发生的异常，我们定义一个 InferenceException 类。

public class InferenceException extends Exception {
    public InferenceException(String message) {
        super(message);
    }

    public InferenceException(String message, Throwable cause) {
        super(message, cause);
    }
}

3.3 定义抽象基类 AbstractInferenceService

为了减少代码重复，我们可以定义一个抽象基类 AbstractInferenceService，该类实现 InferenceService 接口，并提供一些通用的方法。

import java.util.Map;

public abstract class AbstractInferenceService implements InferenceService {

    protected String modelName;
    protected String modelVersion;

    public AbstractInferenceService(String modelName, String modelVersion) {
        this.modelName = modelName;
        this.modelVersion = modelVersion;
    }

    // 可以添加一些通用的辅助方法，例如数据预处理、后处理等。
    protected void preprocessInput(Map<String, Object> input) {
        // 默认实现，可以被子类覆盖
    }

    protected void postprocessOutput(Map<String, Object> output) {
        // 默认实现，可以被子类覆盖
    }

    @Override
    public Map<String, Object> infer(Map<String, Object> input) throws InferenceException {
        preprocessInput(input);
        Map<String, Object> output = doInfer(input);
        postprocessOutput(output);
        return output;
    }

    protected abstract Map<String, Object> doInfer(Map<String, Object> input) throws InferenceException;
}

3.4 实现 TensorFlow Backend

接下来，我们实现一个 TensorFlow Backend，该 Backend 使用 TensorFlow Serving 进行推理。

import org.tensorflow.SavedModelBundle;
import org.tensorflow.Session;
import org.tensorflow.Tensor;
import org.tensorflow.DataType;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Arrays;

public class TensorFlowInferenceService extends AbstractInferenceService {

    private SavedModelBundle model;
    private String modelPath;

    public TensorFlowInferenceService(String modelName, String modelVersion, String modelPath) {
        super(modelName, modelVersion);
        this.modelPath = modelPath;
        loadModel();
    }

    private void loadModel() {
        try {
            model = SavedModelBundle.load(modelPath, "serve");
        } catch (Exception e) {
            throw new RuntimeException("Failed to load TensorFlow model: " + modelPath, e);
        }
    }

    @Override
    protected Map<String, Object> doInfer(Map<String, Object> input) throws InferenceException {
        try {
            Session session = model.session();
            Session.Runner runner = session.runner();

            // 假设输入只有一个名为 "input_tensor" 的 float32 类型的 Tensor
            // 并且输出只有一个名为 "output_tensor" 的 float32 类型的 Tensor
            float[] inputArray = (float[]) input.get("input_tensor");
            int[] shape = new int[]{1, inputArray.length}; // 假设输入是 1xN 的矩阵
            Tensor<Float> inputTensor = Tensor.create(shape, FloatBuffer.wrap(inputArray));

            runner.feed("input_tensor", inputTensor);
            runner.fetch("output_tensor");

            java.util.List<Tensor<?>> results = runner.run();
            Tensor<?> outputTensor = results.get(0);

            float[][] outputArray = new float[1][inputArray.length]; // 假设输出维度与输入相同
            outputTensor.copyTo(outputArray);

            Map<String, Object> output = new HashMap<>();
            output.put("output_tensor", outputArray[0]); // 将float[][]转换为float[]

            inputTensor.close();
            outputTensor.close();

            return output;

        } catch (Exception e) {
            throw new InferenceException("Failed to execute TensorFlow inference", e);
        }
    }

    public static void main(String[] args) throws InferenceException {
        //  TensorFlow 模型路径
        String modelPath = "/path/to/your/tensorflow/model";

        // 创建 TensorFlowInferenceService 实例
        TensorFlowInferenceService inferenceService = new TensorFlowInferenceService("my_model", "1", modelPath);

        // 准备输入数据
        float[] inputData = new float[]{1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        Map<String, Object> input = new HashMap<>();
        input.put("input_tensor", inputData);

        // 执行推理
        Map<String, Object> output = inferenceService.infer(input);

        // 获取推理结果
        float[] outputData = (float[]) output.get("output_tensor");

        // 打印推理结果
        System.out.println("Input: " + Arrays.toString(inputData));
        System.out.println("Output: " + Arrays.toString(outputData));
    }
}

3.5 实现 Triton Backend

接下来，我们实现一个 Triton Backend，该 Backend 使用 Triton Inference Server 进行推理。

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import inference.GrpcServiceGrpc;
import inference.Model;

import java.util.HashMap;
import java.util.Map;

public class TritonInferenceService extends AbstractInferenceService {

    private final String host;
    private final int port;
    private ManagedChannel channel;
    private GrpcServiceGrpc.GrpcServiceBlockingStub blockingStub;

    public TritonInferenceService(String modelName, String modelVersion, String host, int port) {
        super(modelName, modelVersion);
        this.host = host;
        this.port = port;
        init();
    }

    private void init() {
        channel = ManagedChannelBuilder.forAddress(host, port).usePlaintext().build();
        blockingStub = GrpcServiceGrpc.newBlockingStub(channel);
    }

    @Override
    protected Map<String, Object> doInfer(Map<String, Object> input) throws InferenceException {
        try {
            // 构建推理请求
            Model.InferRequest request = Model.InferRequest.newBuilder()
                    .setModelName(modelName)
                    .setModelVersion(modelVersion)
                    // 假设输入只有一个名为 "input_tensor" 的 float32 类型的 Tensor
                    .addInputs(Model.InferRequest.Input.newBuilder()
                            .setName("input_tensor")
                            .setShape(new long[]{1, ((float[])input.get("input_tensor")).length})  // 假设输入是 1xN 的矩阵
                            .setDatatype("FP32")
                            .addContents(Model.InferRequest.InferInputTensorContent.newBuilder().setFp32Contents(convertFloatArrayToList((float[])input.get("input_tensor")))) //需要将float[]转换为List<Float>
                            .build())
                    .build();

            // 执行推理
            Model.InferResponse response = blockingStub.modelInfer(request);

            // 解析推理结果
            Map<String, Object> output = new HashMap<>();
            for (Model.InferResponse.Output outputTensor : response.getOutputsList()) {
                if (outputTensor.getName().equals("output_tensor")) {
                    output.put("output_tensor", convertListToFloatArray(outputTensor.getContents().getFp32ContentsList()));  // 将List<Float>转换为float[]
                }
            }

            return output;

        } catch (Exception e) {
            throw new InferenceException("Failed to execute Triton inference", e);
        }
    }

    private java.util.List<Float> convertFloatArrayToList(float[] array) {
        java.util.List<Float> list = new java.util.ArrayList<>();
        for (float value : array) {
            list.add(value);
        }
        return list;
    }

    private float[] convertListToFloatArray(java.util.List<Float> list) {
        float[] array = new float[list.size()];
        for (int i = 0; i < list.size(); i++) {
            array[i] = list.get(i);
        }
        return array;
    }

    public static void main(String[] args) throws InferenceException {
        // Triton Inference Server 地址
        String host = "localhost";
        int port = 8001;

        // 创建 TritonInferenceService 实例
        TritonInferenceService inferenceService = new TritonInferenceService("my_model", "1", host, port);

        // 准备输入数据
        float[] inputData = new float[]{1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        Map<String, Object> input = new HashMap<>();
        input.put("input_tensor", inputData);

        // 执行推理
        Map<String, Object> output = inferenceService.infer(input);

        // 获取推理结果
        float[] outputData = (float[]) output.get("output_tensor");

        // 打印推理结果
        System.out.println("Input: " + java.util.Arrays.toString(inputData));
        System.out.println("Output: " + java.util.Arrays.toString(outputData));
    }
}

3.6 使用工厂模式创建 InferenceService 实例

为了方便用户选择不同的后端实现，我们可以使用工厂模式来创建 InferenceService 实例。

public class InferenceServiceFactory {

    public static InferenceService create(String backend, Map<String, Object> config) {
        switch (backend) {
            case "tensorflow":
                String modelName = (String) config.get("modelName");
                String modelVersion = (String) config.get("modelVersion");
                String modelPath = (String) config.get("modelPath");
                return new TensorFlowInferenceService(modelName, modelVersion, modelPath);
            case "triton":
                String tritonModelName = (String) config.get("modelName");
                String tritonModelVersion = (String) config.get("modelVersion");
                String host = (String) config.get("host");
                int port = (int) config.get("port");
                return new TritonInferenceService(tritonModelName, tritonModelVersion, host, port);
            case "onnxruntime":
                // TODO: 实现 ONNX Runtime Backend
                throw new IllegalArgumentException("ONNX Runtime backend is not yet implemented.");
            default:
                throw new IllegalArgumentException("Unsupported backend: " + backend);
        }
    }
}

3.7 客户端代码示例

以下代码示例演示了如何在客户端应用程序中使用 InferenceService 接口进行推理。

import java.util.HashMap;
import java.util.Map;
import java.util.Arrays;

public class Client {
    public static void main(String[] args) throws InferenceException {
        // 配置信息
        String backend = "tensorflow"; // 或者 "triton"
        Map<String, Object> config = new HashMap<>();
        if (backend.equals("tensorflow")) {
            config.put("modelName", "my_model");
            config.put("modelVersion", "1");
            config.put("modelPath", "/path/to/your/tensorflow/model");
        } else if (backend.equals("triton")) {
            config.put("modelName", "my_model");
            config.put("modelVersion", "1");
            config.put("host", "localhost");
            config.put("port", 8001);
        }

        // 创建 InferenceService 实例
        InferenceService inferenceService = InferenceServiceFactory.create(backend, config);

        // 准备输入数据
        float[] inputData = new float[]{1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        Map<String, Object> input = new HashMap<>();
        input.put("input_tensor", inputData);

        // 执行推理
        Map<String, Object> output = inferenceService.infer(input);

        // 获取推理结果
        float[] outputData = (float[]) output.get("output_tensor");

        // 打印推理结果
        System.out.println("Input: " + Arrays.toString(inputData));
        System.out.println("Output: " + Arrays.toString(outputData));
    }
}

4. 进一步优化与扩展

• 线程池管理： 对于高并发场景，可以使用线程池来管理推理请求，提高吞吐量。
• 缓存机制： 对于重复的输入数据，可以使用缓存机制来避免重复推理，提高性能。
• 异步推理： 可以使用异步方式执行推理，避免阻塞客户端线程。
• 支持更多数据类型： 可以扩展 InferenceService 接口，支持更多的数据类型，例如图像、文本等。
• 支持更多推理框架： 可以添加对其他推理框架的支持，例如 ONNX Runtime, TensorRT 等。
• 配置管理： 可以使用配置文件来管理后端配置信息，方便用户修改。

5. 表格：不同Backend实现差异

特性	TensorFlow Backend	Triton Backend	ONNX Runtime Backend (示例)
依赖	TensorFlow Java API	gRPC, Triton Inference Server Protobuf	ONNX Runtime Java API
模型格式	SavedModel	Triton Server 支持的多种格式 (TensorFlow, ONNX 等)	ONNX
通信方式	直接加载模型到 JVM	gRPC	直接加载模型到 JVM
输入输出数据处理	需要手动创建 Tensor，处理数据类型转换	需要将 Java 数据类型转换为 Triton 期望的格式	需要手动创建 Tensor，处理数据类型转换
优点	简单易用，适用于单机环境	适用于分布式推理，支持多种模型格式	适用于对延迟敏感的场景，性能较好
缺点	资源占用较高，不适用于大规模分布式推理	需要部署和管理 Triton Inference Server	需要处理不同 ONNX 算子的兼容性问题

6. 总结说明

通过定义统一的 InferenceService 接口和使用工厂模式，我们可以有效地封装跨框架推理接口，适配不同的大模型后端运行环境。这种架构具有良好的可扩展性和灵活性，方便用户根据不同的需求选择不同的后端实现。 此外，根据实际场景进行优化，例如使用线程池、缓存等，可以进一步提高系统的性能和吞吐量。