Java AI 新纪元：利用 Java 16+ 特性构建高性能 AI 应用

// 定义输入样本
public record InputSample(float[] features, String label) {}

// 定义预测输出
public record Prediction(String category, float confidence, Map<String, Float> scores) {}

// 使用
var sample = new InputSample(new float[]{1.2f, 3.4f, 5.6f}, "positive");
var result = new Prediction("cat", 0.92f, Map.of("cat", 0.92f, "dog", 0.08f));

Object rawFeature = getFeature(); // 可能是 float[]、List<Double> 或 double[]
if (rawFeature instanceof float[] arr) {
    // 直接使用 arr，无需强转
    processFloatArray(arr);
} else if (rawFeature instanceof List<?> list && !list.isEmpty() && list.get(0) instanceof Double) {
    // 处理 List<Double>
    float[] converted = list.stream().mapToDouble(d -> (double) d).mapToInt(Double::floatValue).toArray();
    // ...
}

String modelType = "onnx";
int threads = switch (modelType) {
    case "pytorch" -> Runtime.getRuntime().availableProcessors();
    case "tensorflow" -> 2;
    case "onnx" -> 4;
    default -> 1;
};

String prompt = """
    {
      "model": "llama2",
      "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "%s"}
      ],
      "temperature": 0.7
    }
    """.formatted(userInput); // Java 15 引入的 formatted() 方法

public sealed interface AIModel 
    permits ClassificationModel, RegressionModel, GenerativeModel {}

public final class ClassificationModel implements AIModel { ... }
public non-sealed class RegressionModel implements AIModel { ... } // 可被子类扩展

try (Arena arena = Arena.ofConfined()) {
    MemorySegment modelPath = arena.allocateUtf8String("/models/resnet50.onnx");
    // 调用 native 函数创建会话
    SessionHandle session = createSession(modelPath);
    // ...
}

<properties>
    <java.version>17</java.version>
    <djl.version>0.26.0</djl.version>
</properties>

<dependencies>
    <!-- DJL 核心 -->
    <dependency>
        <groupId>ai.djl</groupId>
        <artifactId>api</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <!-- ONNX Runtime 引擎 -->
    <dependency>
        <groupId>ai.djl.onnxruntime</groupId>
        <artifactId>onnxruntime-engine</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <!-- 模型库（用于自动下载预训练模型） -->
    <dependency>
        <groupId>ai.djl</groupId>
        <artifactId>model-zoo</artifactId>
        <version>${djl.version}</version>
    </dependency>
    <!-- Spring Boot Web -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
        <version>3.1.5</version>
    </dependency>
</dependencies>

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.Classifications;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;

import java.nio.file.Path;

public record ImageClassifier(ZooModel<Image, Classifications> model, 
                              Predictor<Image, Classifications> predictor) {

    // 工厂方法，加载 ONNX 模型
    public static ImageClassifier load(Path modelPath, Path synsetPath) throws Exception {
        var criteria = Criteria.builder()
                .optModelPath(modelPath)
                .optEngine("OnnxRuntime")
                .setTypes(Image.class, Classifications.class)
                .optOption("mapLocation", "true") // 允许 GPU
                .optProgress(new ProgressBar())
                .build();

        ZooModel<Image, Classifications> model = ModelZoo.loadModel(criteria);
        Predictor<Image, Classifications> predictor = model.newPredictor();
        return new ImageClassifier(model, predictor);
    }

    public Classifications predict(Image image) throws Exception {
        return predictor.predict(image);
    }

    // 关闭资源
    public void close() throws Exception {
        predictor.close();
        model.close();
    }
}

@RestController
@RequestMapping("/api/v1/classify")
public class ClassifyController {

    private final ImageClassifier classifier;

    public ClassifyController(@Value("${model.path}") String modelPath,
                              @Value("${model.synset}") String synsetPath) throws Exception {
        this.classifier = ImageClassifier.load(Path.of(modelPath), Path.of(synsetPath));
    }

    @PostMapping(consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public PredictionResponse classify(@RequestParam("image") MultipartFile file) throws Exception {
        // 使用 ImageFactory 读取图片
        Image image = ImageFactory.getInstance().fromInputStream(file.getInputStream());

        Classifications result = classifier.predict(image);

        // 提取 Top-5
        var top5 = result.topK(5).stream()
                .map(c -> new ClassEntry(c.getClassName(), c.getProbability()))
                .toList();

        return new PredictionResponse("success", top5);
    }

    // ---- Records 作为 DTO ----
    public record ClassEntry(String label, double confidence) {}

    public record PredictionResponse(String status, List<ClassEntry> predictions) {}
}

// 假设使用 ONNX Runtime 的 Java API
try (Arena arena = Arena.ofConfined()) {
    // 分配连续内存存放 batch 数据
    MemorySegment batchData = arena.allocate(4 * 224 * 224 * 3 * 4); // float32
    // 将多张图片的像素数据写入 batchData（unsafe 方式）
    fillBatchData(batchData, imageList);
    // 调用 native 推理
    float[] results = runInferenceOnBatch(batchData);
}

ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
List<CompletableFuture<PredictionResponse>> futures = requests.stream()
    .map(req -> CompletableFuture.supplyAsync(() -> handleRequest(req), executor))
    .toList();
var allResults = futures.stream().map(CompletableFuture::join).toList();

mvn -Pnative native:compile
./target/image-classifier  # 启动毫秒级