JavaでのMNISTによる文字認識

これは、特定の入力データに対して既存のONNXモデルで推論を実行するための簡単なチュートリアルです。モデルは通常、よく知られたトレーニングフレームワークのいずれかを使用してトレーニングされ、ONNX形式にエクスポートされます。

以下に示すコードは、Java 10以降で利用可能な構文を使用していることに注意してください。Java 8の構文は似ていますが、より冗長です。スコアリングセッションを開始するには、まずOrtEnvironmentを作成し、次にOrtSessionクラスを使用してセッションを開き、モデルへのファイルパスをパラメータとして渡します。

    var env = OrtEnvironment.getEnvironment();
    var session = env.createSession("model.onnx",new OrtSession.SessionOptions());

セッションが作成されると、OrtSessionオブジェクトのrunメソッドを使用してクエリを実行できます。現在、OnnxTensor入力をサポートしており、モデルはOnnxTensor、OnnxSequence、またはOnnxMap出力を生成できます。後者の2つは、scikit-learnなどのフレームワークによって生成されたモデルをスコアリングする場合によく見られます。 run呼び出しはMap<String,OnnxTensor>を想定しており、キーはモデルに格納されている入力ノード名と一致します。これらは、インスタンス化されたセッションでsession.getInputNames()またはsession.getInputInfo()を呼び出すことで表示できます。 run呼び出しはResultオブジェクトを生成します。これには、出力を表すMap<String,OnnxValue>が含まれています。ResultオブジェクトはAutoCloseableであり、try-with-resourcesステートメントで使用して参照がリークするのを防ぐことができます。Resultオブジェクトが閉じられると、そのすべての子OnnxValueも閉じられます。

    OnnxTensor t1,t2;
    var inputs = Map.of("name1",t1,"name2",t2);
    try (var results = session.run(inputs)) {
    // 結果を操作する
   }

入力データをOnnxTensorオブジェクトにいくつかの方法でロードできます。最も効率的な方法はjava.nio.Bufferを使用することですが、多次元配列を使用することも可能です。配列を使用して構築する場合、配列は不規則であってはなりません。

    FloatBuffer sourceData;  // データがFloatBufferにロードされていると仮定します
    long[] dimensions;       // そして入力の次元がここに格納されていると仮定します
    var tensorFromBuffer = OnnxTensor.createTensor(env,sourceData,dimensions);

    float[][] sourceArray = new float[28][28];  // データがfloat配列にロードされていると仮定します
    var tensorFromArray = OnnxTensor.createTensor(env,sourceArray);

これは、事前にトレーニングされたMNISTモデルで推論を実行する完全なサンプルプログラムです。

GPUまたは別のプロバイダーで実行する（オプション）

GPUなどの他の実行プロバイダーを有効にするには、OrtSessionを作成するときにSessionOptionsで適切なフラグをオンにするだけです。

    int gpuDeviceId = 0; // 実行するGPUデバイスID
    var sessionOptions = new OrtSession.SessionOptions();
    sessionOptions.addCUDA(gpuDeviceId);
    var session = environment.createSession("model.onnx", sessionOptions);

実行プロバイダーは、有効化された順に優先されます。