djl 调用onnx模型

调用 ONNX 模型可以使用 DJL 的 ONNX 模块。具体步骤如下：

1. 导入相关依赖

import ai.djl.Model;
import ai.djl.modality.Input;
import ai.djl.modality.Output;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import ai.djl.translate.TranslatorFactory;
import ai.djl.translate.TranslatorFactory.TranslatorFunction;
import ai.djl.translate.TranslatorFactoryImpl;
import ai.djl.translate.TranslatorUtils;
import ai.djl.translate.batch.Batchifier;
import ai.djl.translate.batch.DefaultBatchifier;
import ai.djl.util.Utils;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;

2. 加载 ONNX 模型

Path modelDir = Path.of("path/to/model/dir");
Model model = Model.newInstance("MyModel");
model.setBlock(OnnxModelZoo.resnet50().getBlock());
model.load(modelDir, "model.onnx");

3. 创建 Translator

public class MyTranslator implements Translator<MyInput, MyOutput> {

  private List<String> classes;

  public MyTranslator(List<String> classes) {
    this.classes = classes;
  }

  @Override
  public Batchifier getBatchifier() {
    return DefaultBatchifier.INSTANCE;
  }

  @Override
  public MyOutput processOutput(TranslatorContext ctx, Output output) {
    NDArray array = output.getNDArray();
    return new MyOutput(classes.get(array.argMax().getFloat()));
  }

  @Override
  public MyInput processInput(TranslatorContext ctx, Input input) {
    NDArray array = input.getNDArray();
    return new MyInput(array);
  }
}

TranslatorFactory factory = new TranslatorFactoryImpl();
factory.registerTranslator(MyInput.class, MyOutput.class, new TranslatorFunction<MyInput, MyOutput>() {
  @Override
  public Translator<MyInput, MyOutput> apply(TranslatorContext ctx) {
    List<String> classes = null;
    try {
      classes = Files.readAllLines(Path.of("path/to/classes.txt"));
    } catch (IOException e) {
      e.printStackTrace();
    }
    return new MyTranslator(classes);
  }
});

4. 推理

MyInput input = new MyInput(inputArray);
MyOutput output = model.predict(input).get(0).getOutput(MyOutput.class);

其中，`MyInput` 和 `MyOutput` 分别表示输入和输出的数据类型，需要根据实际情况进行定义。`classes.txt` 文件包含了模型输出的类别信息。