From 47f00b5d49ece70d5905f32a69147812a803f080 Mon Sep 17 00:00:00 2001 From: Adam Pocock Date: Wed, 8 Mar 2023 13:01:08 -0500 Subject: [PATCH] [Java] Initial on device training support (#14027) contributor: @Craigacp --- cmake/onnxruntime_java.cmake | 12 +- cmake/onnxruntime_java_unittests.cmake | 15 +- cmake/onnxruntime_unittests.cmake | 11 +- java/build.gradle | 18 +- .../main/java/ai/onnxruntime/OnnxRuntime.java | 33 +- .../java/ai/onnxruntime/OrtEnvironment.java | 93 ++ .../main/java/ai/onnxruntime/OrtSession.java | 47 +- .../ai/onnxruntime/OrtTrainingSession.java | 964 ++++++++++++++++++ .../onnxruntime/providers/package-info.java | 7 + .../main/native/ai_onnxruntime_OnnxRuntime.c | 14 +- .../ai_onnxruntime_OrtSession_RunOptions.c | 16 + .../ai_onnxruntime_OrtTrainingSession.c | 762 ++++++++++++++ ...me_OrtTrainingSession_OrtCheckpointState.c | 58 ++ .../java/ai/onnxruntime/InferenceTest.java | 4 - .../java/ai/onnxruntime/SparseTensorTest.java | 9 +- .../test/java/ai/onnxruntime/TestHelpers.java | 8 +- .../java/ai/onnxruntime/TrainingTest.java | 256 +++++ .../include/onnxruntime_training_c_api.h | 1 + .../orttraining-linux-gpu-training-apis.yml | 1 + 19 files changed, 2297 insertions(+), 32 deletions(-) create mode 100644 java/src/main/java/ai/onnxruntime/OrtTrainingSession.java create mode 100644 java/src/main/java/ai/onnxruntime/providers/package-info.java create mode 100644 java/src/main/native/ai_onnxruntime_OrtTrainingSession.c create mode 100644 java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c create mode 100644 java/src/test/java/ai/onnxruntime/TrainingTest.java diff --git a/cmake/onnxruntime_java.cmake b/cmake/onnxruntime_java.cmake index d7b47f699a..223f54f60c 100644 --- a/cmake/onnxruntime_java.cmake +++ b/cmake/onnxruntime_java.cmake @@ -1,4 +1,4 @@ -# Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved. +# Copyright (c) 2019, 2022, Oracle and/or its affiliates. All rights reserved. # Licensed under the MIT License. #set(CMAKE_VERBOSE_MAKEFILE on) @@ -57,16 +57,17 @@ file(GLOB onnxruntime4j_native_src "${JAVA_ROOT}/src/main/native/*.c" "${JAVA_ROOT}/src/main/native/*.h" "${REPO_ROOT}/include/onnxruntime/core/session/*.h" + "${REPO_ROOT}/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h" ) # Build the JNI library onnxruntime_add_shared_library_module(onnxruntime4j_jni ${onnxruntime4j_native_src}) -set_property(TARGET onnxruntime4j_jni PROPERTY CXX_STANDARD 11) +set_property(TARGET onnxruntime4j_jni PROPERTY C_STANDARD 11) # depend on java sources. if they change, the JNI should recompile add_dependencies(onnxruntime4j_jni onnxruntime4j) onnxruntime_add_include_to_target(onnxruntime4j_jni onnxruntime_session) # the JNI headers are generated in the onnxruntime4j target -target_include_directories(onnxruntime4j_jni PRIVATE ${REPO_ROOT}/include ${JAVA_ROOT}/build/headers ${JNI_INCLUDE_DIRS}) +target_include_directories(onnxruntime4j_jni PRIVATE ${REPO_ROOT}/include ${REPO_ROOT}/orttraining/orttraining/training_api/include ${JAVA_ROOT}/build/headers ${JNI_INCLUDE_DIRS}) target_link_libraries(onnxruntime4j_jni PUBLIC onnxruntime) set(JAVA_PACKAGE_OUTPUT_DIR ${JAVA_OUTPUT_DIR}/build) @@ -198,7 +199,12 @@ elseif (CMAKE_SYSTEM_NAME STREQUAL "Android") # it is better to not keep a daemon running set(GRADLE_ARGS ${GRADLE_ARGS} --no-daemon) endif() + +# Append relevant native build flags to gradle command set(GRADLE_ARGS ${GRADLE_ARGS} ${ORT_PROVIDER_FLAGS}) +if (onnxruntime_ENABLE_TRAINING_APIS) + set(GRADLE_ARGS ${GRADLE_ARGS} "-DENABLE_TRAINING=1") +endif() message(STATUS "GRADLE_ARGS: ${GRADLE_ARGS}") add_custom_command(TARGET onnxruntime4j_jni POST_BUILD COMMAND ${GRADLE_EXECUTABLE} ${GRADLE_ARGS} WORKING_DIRECTORY ${JAVA_ROOT}) diff --git a/cmake/onnxruntime_java_unittests.cmake b/cmake/onnxruntime_java_unittests.cmake index 7eb326359e..899ec04a71 100644 --- a/cmake/onnxruntime_java_unittests.cmake +++ b/cmake/onnxruntime_java_unittests.cmake @@ -1,4 +1,4 @@ -# Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved. +# Copyright (c) 2019, 2022, Oracle and/or its affiliates. All rights reserved. # Licensed under the MIT License. # This is a windows only file so we can run gradle tests via ctest @@ -6,10 +6,17 @@ FILE(TO_NATIVE_PATH ${GRADLE_EXECUTABLE} GRADLE_NATIVE_PATH) FILE(TO_NATIVE_PATH ${BIN_DIR} BINDIR_NATIVE_PATH) message(STATUS "GRADLE_TEST_EP_FLAGS: ${ORT_PROVIDER_FLAGS}") +if (onnxruntime_ENABLE_TRAINING_APIS) + message(STATUS "Running ORT Java training tests") + execute_process(COMMAND cmd /C ${GRADLE_NATIVE_PATH} --console=plain cmakeCheck -DcmakeBuildDir=${BINDIR_NATIVE_PATH} -Dorg.gradle.daemon=false ${ORT_PROVIDER_FLAGS} -DENABLE_TRAINING=1 + WORKING_DIRECTORY ${REPO_ROOT}/java + RESULT_VARIABLE HAD_ERROR) +else() + execute_process(COMMAND cmd /C ${GRADLE_NATIVE_PATH} --console=plain cmakeCheck -DcmakeBuildDir=${BINDIR_NATIVE_PATH} -Dorg.gradle.daemon=false ${ORT_PROVIDER_FLAGS} + WORKING_DIRECTORY ${REPO_ROOT}/java + RESULT_VARIABLE HAD_ERROR) +endif() -execute_process(COMMAND cmd /C ${GRADLE_NATIVE_PATH} --console=plain cmakeCheck -DcmakeBuildDir=${BINDIR_NATIVE_PATH} -Dorg.gradle.daemon=false ${ORT_PROVIDER_FLAGS} - WORKING_DIRECTORY ${REPO_ROOT}/java - RESULT_VARIABLE HAD_ERROR) if(HAD_ERROR) message(FATAL_ERROR "Java Unitests failed") diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 73b955ede8..e656c5f2a9 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1490,8 +1490,15 @@ if (NOT onnxruntime_BUILD_WEBASSEMBLY) else() add_custom_command(TARGET custom_op_library POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy_if_different $ ${JAVA_NATIVE_TEST_DIR}/$) - add_test(NAME onnxruntime4j_test COMMAND ${GRADLE_EXECUTABLE} cmakeCheck -DcmakeBuildDir=${CMAKE_CURRENT_BINARY_DIR} ${ORT_PROVIDER_FLAGS} - WORKING_DIRECTORY ${REPO_ROOT}/java) + if (onnxruntime_ENABLE_TRAINING_APIS) + message(STATUS "Running Java inference and training tests") + add_test(NAME onnxruntime4j_test COMMAND ${GRADLE_EXECUTABLE} cmakeCheck -DcmakeBuildDir=${CMAKE_CURRENT_BINARY_DIR} ${ORT_PROVIDER_FLAGS} -DENABLE_TRAINING=1 + WORKING_DIRECTORY ${REPO_ROOT}/java) + else() + message(STATUS "Running Java inference tests only") + add_test(NAME onnxruntime4j_test COMMAND ${GRADLE_EXECUTABLE} cmakeCheck -DcmakeBuildDir=${CMAKE_CURRENT_BINARY_DIR} ${ORT_PROVIDER_FLAGS} + WORKING_DIRECTORY ${REPO_ROOT}/java) + endif() endif() set_property(TEST onnxruntime4j_test APPEND PROPERTY DEPENDS onnxruntime4j_jni) endif() diff --git a/java/build.gradle b/java/build.gradle index 2dcd821f3d..1e6c04e2bd 100644 --- a/java/build.gradle +++ b/java/build.gradle @@ -3,7 +3,7 @@ plugins { id 'maven-publish' id 'signing' id 'jacoco' - id 'com.diffplug.spotless' version '5.17.0' + id "com.diffplug.spotless" version "6.13.0" } allprojects { @@ -19,6 +19,7 @@ version = rootProject.file('../VERSION_NUMBER').text.trim() def cmakeBuildDir = System.properties['cmakeBuildDir'] def useCUDA = System.properties['USE_CUDA'] def useROCM = System.properties['USE_ROCM'] +def enableTraining = System.properties['ENABLE_TRAINING'] def cmakeJavaDir = "${cmakeBuildDir}/java" def cmakeNativeLibDir = "${cmakeJavaDir}/native-lib" def cmakeNativeJniDir = "${cmakeJavaDir}/native-jni" @@ -28,7 +29,8 @@ def cmakeBuildOutputDir = "${cmakeJavaDir}/build" def mavenUser = System.properties['mavenUser'] def mavenPwd = System.properties['mavenPwd'] -def mavenArtifactId = (useCUDA != null || useROCM != null) ? project.name + "_gpu" : project.name +def tmpArtifactId = enableTraining == null ? project.name : project.name + "-training" +def mavenArtifactId = (useCUDA == null && useROCM == null) ? tmpArtifactId : tmpArtifactId + "_gpu" java { sourceCompatibility = JavaVersion.VERSION_1_8 @@ -68,6 +70,7 @@ spotless { java { removeUnusedImports() googleJavaFormat() + targetExclude "src/test/java/ai/onnxruntime/OnnxMl.java" } format 'gradle', { target '**/*.gradle' @@ -103,6 +106,7 @@ sourceSets.test { resources.srcDirs += [ "${rootProject.projectDir}/../csharp/testdata", "${rootProject.projectDir}/../onnxruntime/test/testdata", + "${rootProject.projectDir}/../onnxruntime/test/testdata/training_api", "${rootProject.projectDir}/../java/testdata" ] if (cmakeBuildDir != null) { @@ -154,8 +158,8 @@ if (cmakeBuildDir != null) { } dependencies { - testImplementation 'org.junit.jupiter:junit-jupiter-api:5.7.0' - testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.7.0' + testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' + testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' testImplementation 'com.google.protobuf:protobuf-java:3.21.7' } @@ -176,7 +180,7 @@ test { if (cmakeBuildDir != null) { workingDir cmakeBuildDir } - systemProperties System.getProperties().subMap(['USE_CUDA', 'USE_ROCM', 'USE_TENSORRT', 'USE_DNNL', 'USE_OPENVINO', 'JAVA_FULL_TEST']) + systemProperties System.getProperties().subMap(['USE_CUDA', 'USE_ROCM', 'USE_TENSORRT', 'USE_DNNL', 'USE_OPENVINO', 'JAVA_FULL_TEST', 'ENABLE_TRAINING']) testLogging { events "passed", "skipped", "failed" showStandardStreams = true @@ -212,12 +216,12 @@ publishing { } organization { name = 'Microsoft' - url = 'http://www.microsoft.com' + url = 'https://www.microsoft.com' } scm { connection = 'scm:git:git://github.com:microsoft/onnxruntime.git' developerConnection = 'scm:git:ssh://github.com/microsoft/onnxruntime.git' - url = 'http://github.com/microsoft/onnxruntime' + url = 'https://github.com/microsoft/onnxruntime' } developers { developer { diff --git a/java/src/main/java/ai/onnxruntime/OnnxRuntime.java b/java/src/main/java/ai/onnxruntime/OnnxRuntime.java index 430c8fcc58..c58401f3c8 100644 --- a/java/src/main/java/ai/onnxruntime/OnnxRuntime.java +++ b/java/src/main/java/ai/onnxruntime/OnnxRuntime.java @@ -1,5 +1,9 @@ /* +<<<<<<< HEAD * Copyright (c) 2019, 2023, Oracle and/or its affiliates. All rights reserved. +======= + * Copyright (c) 2019, 2022, Oracle and/or its affiliates. All rights reserved. +>>>>>>> 98b8e2e31 (More fixes, tests now pass.) * Licensed under the MIT License. */ package ai.onnxruntime; @@ -38,6 +42,13 @@ final class OnnxRuntime { private static final int ORT_API_VERSION_8 = 8; // Post 1.10 builds of the ORT API private static final int ORT_API_VERSION_11 = 11; + // Post 1.12 builds of the ORT API + private static final int ORT_API_VERSION_13 = 13; + // Post 1.13 builds of the ORT API + private static final int ORT_API_VERSION_14 = 14; + + // The initial release of the ORT training API. + private static final int ORT_TRAINING_API_VERSION_1 = 1; /** * The name of the system property which when set gives the path on disk where the ONNX Runtime @@ -81,6 +92,12 @@ final class OnnxRuntime { /** The API handle. */ static long ortApiHandle; + /** The Training API handle. */ + static long ortTrainingApiHandle; + + /** Is training enabled in the native library */ + static boolean trainingEnabled; + /** The available runtime providers */ static EnumSet providers; @@ -142,10 +159,13 @@ final class OnnxRuntime { load(ONNXRUNTIME_LIBRARY_NAME); load(ONNXRUNTIME_JNI_LIBRARY_NAME); - ortApiHandle = initialiseAPIBase(ORT_API_VERSION_11); + ortApiHandle = initialiseAPIBase(ORT_API_VERSION_14); if (ortApiHandle == 0L) { - throw new IllegalStateException("Failed to load native library"); + throw new IllegalStateException( + "There is a mismatch between the ORT class files and the ORT native library, and the native library could not be loaded"); } + ortTrainingApiHandle = initialiseTrainingAPIBase(ortApiHandle, ORT_API_VERSION_14); + trainingEnabled = ortTrainingApiHandle != 0L; providers = initialiseProviders(ortApiHandle); version = initialiseVersion(); loaded = true; @@ -443,6 +463,15 @@ final class OnnxRuntime { */ private static native long initialiseAPIBase(int apiVersionNumber); + /** + * Get a reference to the training API struct. + * + * @param apiHandle The ORT API struct pointer. + * @param apiVersionNumber The API version to use. + * @return A pointer to the training API struct. + */ + private static native long initialiseTrainingAPIBase(long apiHandle, int apiVersionNumber); + /** * Gets the array of available providers. * diff --git a/java/src/main/java/ai/onnxruntime/OrtEnvironment.java b/java/src/main/java/ai/onnxruntime/OrtEnvironment.java index 73718093c6..22014f44e2 100644 --- a/java/src/main/java/ai/onnxruntime/OrtEnvironment.java +++ b/java/src/main/java/ai/onnxruntime/OrtEnvironment.java @@ -5,8 +5,10 @@ package ai.onnxruntime; import ai.onnxruntime.OrtSession.SessionOptions; +import ai.onnxruntime.OrtTrainingSession.OrtCheckpointState; import java.io.IOException; import java.util.EnumSet; +import java.util.Objects; import java.util.logging.Logger; /** @@ -183,6 +185,15 @@ public final class OrtEnvironment implements AutoCloseable { .addShutdownHook(new Thread(new OrtEnvCloser(OnnxRuntime.ortApiHandle, nativeHandle))); } + /** + * Package accessor for native pointer. + * + * @return The native pointer. + */ + long getNativeHandle() { + return nativeHandle; + } + /** * Create a session using the default {@link SessionOptions}, model and the default memory * allocator. @@ -219,6 +230,7 @@ public final class OrtEnvironment implements AutoCloseable { */ OrtSession createSession(String modelPath, OrtAllocator allocator, SessionOptions options) throws OrtException { + Objects.requireNonNull(modelPath, "model path must not be null"); return new OrtSession(this, modelPath, allocator, options); } @@ -258,9 +270,90 @@ public final class OrtEnvironment implements AutoCloseable { */ OrtSession createSession(byte[] modelArray, OrtAllocator allocator, SessionOptions options) throws OrtException { + Objects.requireNonNull(modelArray, "model array must not be null"); return new OrtSession(this, modelArray, allocator, options); } + /** + * Create a training session using the default {@link SessionOptions}, model and the default + * memory allocator. + * + * @param checkpointPath Path to the checkpoint folder. + * @param trainPath Path to the training model. + * @param evalPath Path to the evaluation model. Null signifies there is no eval model. + * @param optimizerPath Path to the optimizer model. Null signifies there is no optimizer model. + * @return An {@link OrtTrainingSession} with the specified model loaded. + * @throws OrtException If the model failed to load, wasn't compatible or caused an error. + */ + public OrtTrainingSession createTrainingSession( + String checkpointPath, String trainPath, String evalPath, String optimizerPath) + throws OrtException { + return createTrainingSession( + checkpointPath, trainPath, evalPath, optimizerPath, new OrtSession.SessionOptions()); + } + + /** + * Create a training session using the specified {@link SessionOptions}, model and the default + * memory allocator. + * + * @param checkpointPath Path to the checkpoint folder. + * @param trainPath Path to the training model. + * @param evalPath Path to the evaluation model. Null signifies there is no eval model. + * @param optimizerPath Path to the optimizer model. Null signifies there is no optimizer model. + * @param options The session options. + * @return An {@link OrtTrainingSession} with the specified model. + * @throws OrtException If the model failed to load, wasn't compatible or caused an error. + */ + public OrtTrainingSession createTrainingSession( + String checkpointPath, + String trainPath, + String evalPath, + String optimizerPath, + SessionOptions options) + throws OrtException { + return createTrainingSession( + checkpointPath, trainPath, evalPath, optimizerPath, defaultAllocator, options); + } + + /** + * Create a training session using the specified {@link SessionOptions} and model. + * + * @param checkpointPath Path to the checkpoint folder. + * @param trainPath Path to the training model. + * @param evalPath Path to the evaluation model. + * @param optimizerPath Path to the optimizer model. + * @param allocator The memory allocator to use. + * @param options The session options. + * @return An {@link OrtTrainingSession} with the specified model. + * @throws OrtException If the model failed to load, wasn't compatible or caused an error. + */ + OrtTrainingSession createTrainingSession( + String checkpointPath, + String trainPath, + String evalPath, + String optimizerPath, + OrtAllocator allocator, + SessionOptions options) + throws OrtException { + if (OnnxRuntime.trainingEnabled) { + Objects.requireNonNull(trainPath, "train path must not be null"); + OrtCheckpointState checkpointState = OrtCheckpointState.loadCheckpoint(checkpointPath); + return new OrtTrainingSession( + this, allocator, options, checkpointState, trainPath, evalPath, optimizerPath); + } else { + throw new IllegalStateException("Training is not enabled in this build of ONNX Runtime."); + } + } + + /** + * Is training enabled in this build of ONNX Runtime? + * + * @return True if training is enabled. + */ + public boolean isTrainingEnabled() { + return OnnxRuntime.trainingEnabled; + } + /** * Turns on or off the telemetry. * diff --git a/java/src/main/java/ai/onnxruntime/OrtSession.java b/java/src/main/java/ai/onnxruntime/OrtSession.java index df6f27741f..536fd99598 100644 --- a/java/src/main/java/ai/onnxruntime/OrtSession.java +++ b/java/src/main/java/ai/onnxruntime/OrtSession.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019, 2022, Oracle and/or its affiliates. All rights reserved. + * Copyright (c) 2019, 2023, Oracle and/or its affiliates. All rights reserved. * Licensed under the MIT License. */ package ai.onnxruntime; @@ -19,6 +19,7 @@ import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.logging.Logger; @@ -70,7 +71,8 @@ public class OrtSession implements AutoCloseable { OrtSession(OrtEnvironment env, String modelPath, OrtAllocator allocator, SessionOptions options) throws OrtException { this( - createSession(OnnxRuntime.ortApiHandle, env.nativeHandle, modelPath, options.nativeHandle), + createSession( + OnnxRuntime.ortApiHandle, env.getNativeHandle(), modelPath, options.getNativeHandle()), allocator); } @@ -86,7 +88,8 @@ public class OrtSession implements AutoCloseable { OrtSession(OrtEnvironment env, byte[] modelArray, OrtAllocator allocator, SessionOptions options) throws OrtException { this( - createSession(OnnxRuntime.ortApiHandle, env.nativeHandle, modelArray, options.nativeHandle), + createSession( + OnnxRuntime.ortApiHandle, env.getNativeHandle(), modelArray, options.getNativeHandle()), allocator); } @@ -292,7 +295,7 @@ public class OrtSession implements AutoCloseable { "Unknown output name " + s + ", expected one of " + outputNames.toString()); } } - long runOptionsHandle = runOptions == null ? 0 : runOptions.nativeHandle; + long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle(); OnnxValue[] outputValues = run( @@ -554,6 +557,15 @@ public class OrtSession implements AutoCloseable { } } + /** + * Package accessor for the native pointer. + * + * @return The native pointer. + */ + long getNativeHandle() { + return nativeHandle; + } + /** * Sets the execution mode of this options object, overriding the old setting. * @@ -697,6 +709,7 @@ public class OrtSession implements AutoCloseable { */ public void registerCustomOpLibrary(String path) throws OrtException { checkClosed(); + Objects.requireNonNull(path, "path must not be null"); long customHandle = registerCustomOpLibrary(OnnxRuntime.ortApiHandle, nativeHandle, path); customLibraryHandles.add(customHandle); } @@ -1164,6 +1177,15 @@ public class OrtSession implements AutoCloseable { this.nativeHandle = createRunOptions(OnnxRuntime.ortApiHandle); } + /** + * Package accessor for native pointer. + * + * @return The native pointer. + */ + long getNativeHandle() { + return nativeHandle; + } + /** * Sets the current logging level on this RunOptions. * @@ -1243,6 +1265,20 @@ public class OrtSession implements AutoCloseable { setTerminate(OnnxRuntime.ortApiHandle, nativeHandle, terminate); } + /** + * Adds a configuration entry to this {@code RunOptions}. + * + *

Setting the same key will overwrite the value. + * + * @param key The configuration key. + * @param value The configuration value. + * @throws OrtException If the native library call failed. + */ + public void addRunConfigEntry(String key, String value) throws OrtException { + checkClosed(); + addRunConfigEntry(OnnxRuntime.ortApiHandle, nativeHandle, key, value); + } + /** Checks if the RunOptions is closed, if so throws {@link IllegalStateException}. */ private void checkClosed() { if (closed) { @@ -1280,6 +1316,9 @@ public class OrtSession implements AutoCloseable { private native void setTerminate(long apiHandle, long nativeHandle, boolean terminate) throws OrtException; + private native void addRunConfigEntry( + long apiHandle, long nativeHandle, String key, String value) throws OrtException; + private static native void close(long apiHandle, long nativeHandle); } diff --git a/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java b/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java new file mode 100644 index 0000000000..58f27e47cd --- /dev/null +++ b/java/src/main/java/ai/onnxruntime/OrtTrainingSession.java @@ -0,0 +1,964 @@ +/* + * Copyright (c) 2022, 2023, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +package ai.onnxruntime; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +/** + * Wraps an ONNX training model and allows training and inference calls. + * + *

Allows the inspection of the model's input and output nodes. Produced by an {@link + * OrtEnvironment}. + * + *

Most instance methods throw {@link IllegalStateException} if the session is closed and the + * methods are called. + */ +public final class OrtTrainingSession implements AutoCloseable { + + static { + try { + OnnxRuntime.init(); + } catch (IOException e) { + throw new RuntimeException("Failed to load onnx-runtime library", e); + } + } + + private final long nativeHandle; + private final OrtAllocator allocator; + private final OrtCheckpointState checkpoint; + + private final String trainPath; + private final String evalPath; + private final String optimizerPath; + + private final Set trainInputNames; + private final Set trainOutputNames; + + private final Set evalInputNames; + private final Set evalOutputNames; + + private boolean closed = false; + + /** + * Constructs an {@code OrtTrainingSession}. + * + *

Note the guard on training being enabled is not present in this method, and it should only + * be called after {@link OnnxRuntime#trainingEnabled} has been checked to be true. + * + * @param env The environment. + * @param allocator The memory allocator. + * @param options The session options. + * @param checkpoint The checkpoint to load. + * @param trainPath The path to the training model. + * @param evalPath The path to the evaluation model. + * @param optimizerPath The path to the optimizer model. + * @throws OrtException If the native creation failed. + */ + OrtTrainingSession( + OrtEnvironment env, + OrtAllocator allocator, + OrtSession.SessionOptions options, + OrtCheckpointState checkpoint, + String trainPath, + String evalPath, + String optimizerPath) + throws OrtException { + this( + createTrainingSession( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + env.getNativeHandle(), + options.getNativeHandle(), + checkpoint.nativeHandle, + trainPath, + evalPath, + optimizerPath), + allocator, + checkpoint, + trainPath, + evalPath, + optimizerPath); + } + + /** + * Wraps an OrtTrainingSession around the native session pointer. + * + * @param nativeHandle The native session pointer. + * @param allocator The memory allocator. + * @param trainPath The path on disk to the training model. + * @param evalPath The path on disk to the evaluation model. + * @param optimizerPath The path on disk to the optimizer model. + */ + private OrtTrainingSession( + long nativeHandle, + OrtAllocator allocator, + OrtCheckpointState checkpoint, + String trainPath, + String evalPath, + String optimizerPath) + throws OrtException { + this.nativeHandle = nativeHandle; + this.allocator = allocator; + this.checkpoint = checkpoint; + this.trainPath = trainPath; + this.evalPath = evalPath; + this.optimizerPath = optimizerPath; + + this.trainInputNames = + Collections.unmodifiableSet( + new LinkedHashSet<>( + Arrays.asList( + getTrainInputNames( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle)))); + this.trainOutputNames = + Collections.unmodifiableSet( + new LinkedHashSet<>( + Arrays.asList( + getTrainOutputNames( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle)))); + this.evalInputNames = + Collections.unmodifiableSet( + new LinkedHashSet<>( + Arrays.asList( + getEvalInputNames( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle)))); + this.evalOutputNames = + Collections.unmodifiableSet( + new LinkedHashSet<>( + Arrays.asList( + getEvalOutputNames( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle)))); + } + + /* + * \brief Create a training session that can be used to begin or resume training. + * + *

This function creates a training session based on the env and session options provided that + * can begin or resume training from a given checkpoint state for the given onnx models. The + * checkpoint state represents the parameters of the training session which will be moved to the + * device specified by the user through the session options (if necessary). + * + *

\param[in] env Environment to be used for the training session. \param[in] options Session + * options that the user can customize for this training session. \param[in] checkpoint_state + * Training states that the training session uses as a starting point for training. \param[in] + * train_model_path Model to be used to perform training that can be generated using the offline + * tooling library. \param[in] eval_model_path Model to be used to perform evaluation that can be + * generated using the offline tooling library. \param[in] optimizer_model_path Model to be used + * to the optimizer step for weight updates. The model can be generated using the offline tooling + * library. \param[out] out Created training session. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* + * options, _Inout_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* train_model_path, + * _In_ const ORTCHAR_T* eval_model_path, _In_ const ORTCHAR_T* optimizer_model_path, _Outptr_ + * OrtTrainingSession** out); + */ + private static native long createTrainingSession( + long apiHandle, + long trainingHandle, + long envHandle, + long optionsHandle, + long checkpointHandle, + String trainPath, + String evalPath, + String optimizerPath); + + /** + * Returns an ordered set of the train model input names. + * + * @return The training inputs. + */ + public Set getTrainInputNames() { + return trainInputNames; + } + + /** + * Returns an ordered set of the train model output names. + * + * @return The training outputs. + */ + public Set getTrainOutputNames() { + return trainOutputNames; + } + + /** + * Returns an ordered set of the eval model input names. + * + * @return The evaluation inputs. + */ + public Set getEvalInputNames() { + return evalInputNames; + } + + /** + * Returns an ordered set of the eval model output names. + * + * @return The evaluation outputs. + */ + public Set getEvalOutputNames() { + return evalOutputNames; + } + + /** Checks if the OrtTrainingSession is closed, if so throws {@link IllegalStateException}. */ + private void checkClosed() { + if (closed) { + throw new IllegalStateException("Trying to use a closed OrtTrainingSession"); + } + } + + @Override + public void close() { + if (!closed) { + closeSession(OnnxRuntime.ortTrainingApiHandle, nativeHandle); + checkpoint.close(); + closed = true; + } else { + throw new IllegalStateException("Trying to close an already closed OrtSession."); + } + } + + private native void closeSession(long trainingHandle, long nativeHandle); + + /** + * Save out the training session state into the supplied checkpoint directory. + * + * @param outputPath Path to a checkpoint directory. + * @param saveOptimizer Should the optimizer states be saved out. + * @throws OrtException If the native call failed. + */ + public void saveCheckpoint(Path outputPath, boolean saveOptimizer) throws OrtException { + checkClosed(); + String outputStr = outputPath.toString(); + saveCheckpoint( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + outputStr, + saveOptimizer); + } + + /* + * \brief Save the training session states to a checkpoint directory on disk. + * + *

This function retrieves the training session states from the training session and serializes + * them to a checkpoint directory on disk. This checkpoint can later be loaded by invoking + * LoadCheckpoint to continue the training with the same states. + * + *

\param[in] checkpoint_path Path to the checkpoint directory \param[in] session The training + * session from where the checkpoint states are to be retrieved. \param[in] save_optimizer_state + * Boolean flag indicating whether or not to save the optimizer states to the checkpoint. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _In_ const + * OrtTrainingSession* session, bool save_optimizer_state); + */ + private native void saveCheckpoint( + long apiHandle, long trainingHandle, long nativeHandle, String path, boolean saveOptimizer) + throws OrtException; + + /* + * \brief Retrieves the number of user outputs in the training model. + * + *

This function returns the number of outputs of the training model so that the user can + * allocate space for the number of outputs when TrainStep is invoked. + * + *

\param[in] sess The training session which has working knowledge of the training model. + * \param[out] out Number of user outputs in the training model. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputCount, _In_ const OrtTrainingSession* + * sess, _Out_ size_t* out); ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputName, _In_ const + * OrtSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output); + */ + private native String[] getTrainInputNames( + long apiHandle, long trainingApiHandle, long nativeHandle, long allocatorHandle) + throws OrtException; + + private native String[] getTrainOutputNames( + long apiHandle, long trainingApiHandle, long nativeHandle, long allocatorHandle) + throws OrtException; + + /* + * \brief Retrieves the number of user outputs in the eval model. + * + *

This function returns the number of outputs of the eval model so that the user can allocate + * space for the number of outputs when EvalStep is invoked. + * + *

\param[in] sess The training session which has working knowledge of the eval model. + * \param[out] out Number of user outputs in the eval model. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(TrainingSessionGetEvalModelOutputCount, _In_ const OrtTrainingSession* sess, + * _Out_ size_t* out); ORT_API2_STATUS(TrainingSessionGetEvalModelOutputName, _In_ const + * OrtSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output); + */ + private native String[] getEvalInputNames( + long apiHandle, long trainingApiHandle, long nativeHandle, long allocatorHandle) + throws OrtException; + + private native String[] getEvalOutputNames( + long apiHandle, long trainingApiHandle, long nativeHandle, long allocatorHandle) + throws OrtException; + + /** + * Ensures the gradients are reset to zero before the next call to {@link #trainStep}. + * + *

Note this is a lazy call, the gradients are cleared as part of running the next {@link + * #trainStep} and not before. + * + * @throws OrtException If the native call failed. + */ + public void lazyResetGrad() throws OrtException { + checkClosed(); + lazyResetGrad(OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, nativeHandle); + } + + /* + * \brief Reset the training model gradients to zero lazily. + * + *

This function sets the internal state of the training session such that the training model + * gradients will be reset just before the new gradients are computed on the next invocation of + * TrainStep. + * + *

\param[in] session The training session which has working knowledge of the eval model. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(ResetGrad, _Inout_ OrtTrainingSession* session); + */ + private native void lazyResetGrad(long apiHandle, long trainingHandle, long nativeHandle) + throws OrtException; + + /** + * Sets the RNG seed used by ONNX Runtime. + * + *

Note this setting is global across OrtTrainingSession instances. + * + * @param seed The RNG seed. + * @throws OrtException If the native call failed. + */ + public static void setSeed(long seed) throws OrtException { + setSeed(OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, seed); + } + + private static native void setSeed(long apiHandle, long trainingHandle, long seed) + throws OrtException; + + /** + * Performs a single step of training, accumulating the gradients. + * + * @param inputs The inputs (must include both the features and the target). + * @return All outputs produced by the training step. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result trainStep(Map inputs) + throws OrtException { + return trainStep(inputs, trainOutputNames, null); + } + + /** + * Performs a single step of training, accumulating the gradients. + * + * @param inputs The inputs (must include both the features and the target). + * @param runOptions Run options for controlling this specific call. + * @return All outputs produced by the training step. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result trainStep( + Map inputs, OrtSession.RunOptions runOptions) + throws OrtException { + return trainStep(inputs, trainOutputNames, runOptions); + } + + /** + * Performs a single step of training, accumulating the gradients. + * + * @param inputs The inputs (must include both the features and the target). + * @param requestedOutputs The requested outputs. + * @return Requested outputs produced by the training step. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result trainStep( + Map inputs, Set requestedOutputs) + throws OrtException { + return trainStep(inputs, requestedOutputs, null); + } + + /** + * Performs a single step of training, accumulating the gradients. + * + * @param inputs The inputs (must include both the features and the target). + * @param requestedOutputs The requested outputs. + * @param runOptions Run options for controlling this specific call. + * @return Requested outputs produced by the training step. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result trainStep( + Map inputs, + Set requestedOutputs, + OrtSession.RunOptions runOptions) + throws OrtException { + checkClosed(); + if ((inputs.isEmpty() && (trainInputNames.size() != 0)) + || (inputs.size() > trainInputNames.size())) { + throw new OrtException( + "Unexpected number of inputs, expected [1," + + trainInputNames.size() + + ") found " + + inputs.size()); + } + if (requestedOutputs.isEmpty() || (requestedOutputs.size() > trainOutputNames.size())) { + throw new OrtException( + "Unexpected number of requestedOutputs, expected [1," + + trainOutputNames.size() + + ") found " + + requestedOutputs.size()); + } + String[] inputNamesArray = new String[inputs.size()]; + long[] inputHandles = new long[inputs.size()]; + int i = 0; + for (Map.Entry t : inputs.entrySet()) { + if (trainInputNames.contains(t.getKey())) { + inputNamesArray[i] = t.getKey(); + inputHandles[i] = t.getValue().getNativeHandle(); + i++; + } else { + throw new OrtException( + "Unknown input name " + t.getKey() + ", expected one of " + trainInputNames); + } + } + String[] outputNamesArray = new String[requestedOutputs.size()]; + i = 0; + for (String s : requestedOutputs) { + if (trainOutputNames.contains(s)) { + outputNamesArray[i] = s; + i++; + } else { + throw new OrtException( + "Unknown output name " + s + ", expected one of " + trainOutputNames.toString()); + } + } + long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle(); + + OnnxValue[] outputValues = + trainStep( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle, + inputNamesArray, + inputHandles, + inputNamesArray.length, + outputNamesArray, + outputNamesArray.length, + runOptionsHandle); + return new OrtSession.Result(outputNamesArray, outputValues); + } + + /* + * \brief Computes the outputs and the gradients for the training model for the given inputs + * + *

This function performs a training step that computes the outputs and the gradients of the + * training model for the given inputs. The train step is performed based on the training model + * that was provided to the training session. The gradients computed are stored inside the + * training session so they can be later consumed by the OptimizerStep function. + * + *

\param[in] sess The training session which has working knowledge of the eval model. + * \param[in] run_options Run options for this training step. \param[in] inputs_len Number of user + * inputs to the training model. \param[in] inputs The user inputs to the training model. + * \param[in] outputs_len Number of user outputs expected from this training step. \param[out] + * outputs User outputs computed by train step. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* + * run_options, size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, size_t + * outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + */ + private native OnnxValue[] trainStep( + long apiHandle, + long trainingApiHandle, + long nativeHandle, + long allocatorHandle, + String[] inputNamesArray, + long[] inputs, + long numInputs, + String[] outputNamesArray, + long numOutputs, + long runOptionsHandle); + + /** + * Performs a single evaluation step using the supplied inputs. + * + * @param inputs The model inputs. + * @return All model outputs. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result evalStep(Map inputs) + throws OrtException { + return evalStep(inputs, evalOutputNames, null); + } + + /** + * Performs a single evaluation step using the supplied inputs. + * + * @param inputs The model inputs. + * @param runOptions Run options for controlling this specific call. + * @return All model outputs. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result evalStep( + Map inputs, OrtSession.RunOptions runOptions) + throws OrtException { + return evalStep(inputs, evalOutputNames, runOptions); + } + + /** + * Performs a single evaluation step using the supplied inputs. + * + * @param inputs The model inputs. + * @param requestedOutputs The requested output names. + * @return The requested outputs. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result evalStep( + Map inputs, Set requestedOutputs) + throws OrtException { + return evalStep(inputs, requestedOutputs, null); + } + + /** + * Performs a single evaluation step using the supplied inputs. + * + * @param inputs The model inputs. + * @param requestedOutputs The requested output names. + * @param runOptions Run options for controlling this specific call. + * @return The requested outputs. + * @throws OrtException If the native call failed. + */ + public OrtSession.Result evalStep( + Map inputs, + Set requestedOutputs, + OrtSession.RunOptions runOptions) + throws OrtException { + checkClosed(); + if ((inputs.isEmpty() && (evalInputNames.size() != 0)) + || (inputs.size() > evalInputNames.size())) { + throw new OrtException( + "Unexpected number of inputs, expected [1," + + evalInputNames.size() + + ") found " + + inputs.size()); + } + if (requestedOutputs.isEmpty() || (requestedOutputs.size() > evalOutputNames.size())) { + throw new OrtException( + "Unexpected number of requestedOutputs, expected [1," + + evalOutputNames.size() + + ") found " + + requestedOutputs.size()); + } + String[] inputNamesArray = new String[inputs.size()]; + long[] inputHandles = new long[inputs.size()]; + int i = 0; + for (Map.Entry t : inputs.entrySet()) { + if (evalInputNames.contains(t.getKey())) { + inputNamesArray[i] = t.getKey(); + inputHandles[i] = t.getValue().getNativeHandle(); + i++; + } else { + throw new OrtException( + "Unknown input name " + t.getKey() + ", expected one of " + evalInputNames.toString()); + } + } + String[] outputNamesArray = new String[requestedOutputs.size()]; + i = 0; + for (String s : requestedOutputs) { + if (evalOutputNames.contains(s)) { + outputNamesArray[i] = s; + i++; + } else { + throw new OrtException( + "Unknown output name " + s + ", expected one of " + evalOutputNames.toString()); + } + } + long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle(); + + OnnxValue[] outputValues = + evalStep( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + allocator.handle, + inputNamesArray, + inputHandles, + inputNamesArray.length, + outputNamesArray, + outputNamesArray.length, + runOptionsHandle); + return new OrtSession.Result(outputNamesArray, outputValues); + } + + /* + * \brief Computes the outputs for the eval model for the given inputs + * + *

This function performs an eval step that computes the outputs of the eval model for the + * given inputs. The eval step is performed based on the eval model that was provided to the + * training session. + * + *

\param[in] sess The training session which has working knowledge of the eval model. + * \param[in] run_options Run options for this eval step. \param[in] inputs_len Number of user + * inputs to the eval model. \param[in] inputs The user inputs to the eval model. \param[in] + * outputs_len Number of user outputs expected from this eval step. \param[out] outputs User + * outputs computed by eval step. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(EvalStep, _In_ const OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* + * run_options, size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, size_t + * outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + */ + private native OnnxValue[] evalStep( + long apiHandle, + long trainingApiHandle, + long nativeHandle, + long allocatorHandle, + String[] inputNamesArray, + long[] inputs, + long numInputs, + String[] outputNamesArray, + long numOutputs, + long runOptionsHandle) + throws OrtException; + + /** + * Sets the learning rate for the training session. + * + *

Should be used only when there is no learning rate scheduler in the session. Not used to set + * the initial learning rate for LR schedulers. + * + * @param learningRate The learning rate. + * @throws OrtException If the call failed. + */ + public void setLearningRate(float learningRate) throws OrtException { + checkClosed(); + setLearningRate( + OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, nativeHandle, learningRate); + } + + /* + * \brief Sets the learning rate for this training session. + * + *

This function allows users to set the learning rate for the training session. The current + * learning rate is maintained by the training session and can be overwritten by invoking this + * function with the desired learning rate. This function should not be used when a valid learning + * rate scheduler is registered. It should be used either to set the learning rate derived from a + * custom learning rate scheduler or to set the learning rate constant to be used throughout the + * training session. Please note that this function does not set the initial learning rate that + * may be needed by the predefined learning rate schedulers. To set the initial learning rate for + * learning rate schedulers, please look at the function `RegisterLinearLRScheduler`. + * + *

\param[in] sess The training session on which the learning rate needs to be set. \param[in] + * learning_rate Desired learning rate to set. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(SetLearningRate, _Inout_ OrtTrainingSession* sess, _In_ float + * learning_rate); + */ + private native void setLearningRate( + long apiHandle, long trainingApiHandle, long nativeHandle, float learningRate) + throws OrtException; + + /** + * Gets the current learning rate for this training session. + * + * @return The current learning rate. + * @throws OrtException If the call failed. + */ + public float getLearningRate() throws OrtException { + checkClosed(); + return getLearningRate( + OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, nativeHandle); + } + + /* + * \brief Gets the current learning rate for this training session. + * + *

This function allows users to get the learning rate for the training session. The current + * learning rate is maintained by the training session + * + *

\param[in] sess The training session on which the learning rate needs to be set. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(GetLearningRate, _Inout_ OrtTrainingSession* sess, _Out_ float* + * learning_rate); + */ + private native float getLearningRate(long apiHandle, long trainingApiHandle, long nativeHandle); + + /** + * Applies the gradient updates to the trainable parameters using the optimizer model. + * + * @throws OrtException If the native call failed. + */ + public void optimizerStep() throws OrtException { + optimizerStep(null); + } + + /** + * Applies the gradient updates to the trainable parameters using the optimizer model. + * + *

The run options can be used to control logging and to terminate the call early. + * + * @param runOptions Options for controlling the model execution. + * @throws OrtException If the native call failed. + */ + public void optimizerStep(OrtSession.RunOptions runOptions) throws OrtException { + checkClosed(); + long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle(); + optimizerStep( + OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, nativeHandle, runOptionsHandle); + } + + /* + * \brief Performs the weight updates for the trainable parameters using the optimizer model. + * + *

This function performs the weight update step that updates the trainable parameters such + * that they take a step in the direction of their gradients. The optimizer step is performed + * based on the optimizer model that was provided to the training session. The updated parameters + * are stored inside the training session so that they can be used by the next TrainStep function + * call. + * + *

\param[in] sess The training session which has working knowledge of the optimizer model. + * \param[in] run_options Run options for this eval step. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(OptimizerStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const + * OrtRunOptions* run_options); + */ + private native void optimizerStep( + long apiHandle, long trainingApiHandle, long nativeHandle, long runOptionsHandle) + throws OrtException; + + /** + * Registers a linear learning rate scheduler with linear warmup. + * + * @param warmupSteps The number of steps to increase the learning rate from zero to {@code + * initialLearningRate}. + * @param totalSteps The total number of steps this scheduler operates over. + * @param initialLearningRate The maximum learning rate. + * @throws OrtException If the native call failed. + */ + public void registerLinearLRScheduler( + long warmupSteps, long totalSteps, float initialLearningRate) throws OrtException { + registerLinearLRScheduler( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + warmupSteps, + totalSteps, + initialLearningRate); + } + + /* + * \brief Registers the use of the Linear learning rate scheduler for the training session. + * + *

Register a Linear learning rate scheduler with the given learning rate scheduler parameters. + * Specify the initial learning rate that should be used with this learning rate scheduler and + * training session. + * + *

\param[in] sess The training session that should use the linear learning rate scheduler. + * \param[in] warmup_step_count Warmup steps for LR warmup. \param[in] total_step_count Total step + * count. \param[in] initial_lr The initial learning rate to be used by the training session. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(RegisterLinearLRScheduler, _Inout_ OrtTrainingSession* sess, _In_ const + * int64_t warmup_step_count, _In_ const int64_t total_step_count, _In_ const float initial_lr); + */ + private native void registerLinearLRScheduler( + long apiHandle, + long trainingApiHandle, + long nativeHandle, + long warmupSteps, + long totalSteps, + float initialLearningRate) + throws OrtException; + + /** + * Updates the learning rate based on the registered learning rate scheduler. + * + * @throws OrtException If the native call failed. + */ + public void schedulerStep() throws OrtException { + checkClosed(); + schedulerStep(OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, nativeHandle); + } + + /* + * \brief Update the learning rate based on the registered learing rate scheduler. + * + *

Takes a scheduler step that updates the learning rate that is being used by the training + * session. This function should typically be called before invoking the optimizer step for each + * round, or as determined necessary to update the learning rate being used by the training + * session. Please note that a valid predefined learning rate scheduler must be first registered + * to invoke this function. + * + *

\param[in] sess The training session that has the registered learning rate scheduler. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(SchedulerStep, _Inout_ OrtTrainingSession* sess); + */ + private native void schedulerStep(long apiHandle, long trainingApiHandle, long nativeHandle) + throws OrtException; + + /** + * Exports the evaluation model as a model suitable for inference, setting the desired nodes as + * output nodes. + * + *

Note that this method reloads the evaluation model from the path provided to the training + * session, and this path must still be valid. + * + * @param outputPath The path to write out the inference model. + * @param outputNames The names of the output nodes. + * @throws OrtException If the native call failed. + */ + public void exportModelForInference(Path outputPath, String[] outputNames) throws OrtException { + checkClosed(); + if (outputNames.length == 0) { + throw new IllegalArgumentException("Requires at least one output name"); + } + String outputStr = outputPath.toString(); + exportModelForInference( + OnnxRuntime.ortApiHandle, + OnnxRuntime.ortTrainingApiHandle, + nativeHandle, + outputStr, + outputNames.length, + outputNames); + } + + /* + * \brief Export a model that can be used for inferencing. + * + *

If the training session was provided with an eval model, the training session can generate + * an inference model if it knows the inference graph outputs. The input inference graph outputs + * are used to prune the eval model so that the output model's outputs align with the provided + * outputs. The exported model is saved at the path provided and can be used for inferencing with + * InferenceSession. Note that the function re-loads the eval model from the path provided to + * CreateTrainingSession and expects that this path still be valid. + * + *

\param[in] sess The training session. \param[in] inference_model_path Path where the + * inference model should be serialized to. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(ExportModelForInferencing, _Inout_ OrtTrainingSession* sess, _In_ const + * ORTCHAR_T* inference_model_path, size_t graph_outputs_len, _In_reads_(graph_outputs_len) const + * char* const* graph_output_names); + */ + private native void exportModelForInference( + long apiHandle, + long trainingApiHandle, + long nativeHandle, + String outputPath, + long numOutputs, + String[] outputNames) + throws OrtException; + + /** Wrapper class for the checkpoint state. */ + static final class OrtCheckpointState implements AutoCloseable { + final long nativeHandle; + + /** + * Wraps an object around the checkpoint native handle. + * + * @param nativeHandle The pointer to the checkpoint. + */ + OrtCheckpointState(long nativeHandle) { + this.nativeHandle = nativeHandle; + } + + /** + * Loads a checkpoint from disk. + * + * @param checkpointPath The path to load + * @return The checkpoint. + * @throws OrtException If the checkpoint failed to load. + */ + static OrtCheckpointState loadCheckpoint(Path checkpointPath) throws OrtException { + String pathStr = checkpointPath.toString(); + return loadCheckpoint(pathStr); + } + + /** + * Loads a checkpoint from disk. + * + * @param checkpoint The path to load + * @return The checkpoint. + * @throws OrtException If the checkpoint failed to load. + */ + static OrtCheckpointState loadCheckpoint(String checkpoint) throws OrtException { + if (OnnxRuntime.trainingEnabled) { + Objects.requireNonNull(checkpoint, "checkpoint path must not be null"); + return new OrtCheckpointState( + loadCheckpoint(OnnxRuntime.ortApiHandle, OnnxRuntime.ortTrainingApiHandle, checkpoint)); + } else { + throw new IllegalStateException("Training is not enabled in this build of ONNX Runtime."); + } + } + + @Override + public void close() { + close(OnnxRuntime.ortTrainingApiHandle, nativeHandle); + } + + /* + * \brief Load a checkpoint state from directory on disk into checkpoint_state. + * + *

This function will parse a checkpoint directory, pull relevant files and load the training + * states into the checkpoint_state. This checkpoint state can then be used to create the + * training session by invoking CreateTrainingSession. By doing so, the training session will + * resume training from the given checkpoint. + * + *

\param[in] checkpoint_path Path to the checkpoint directory \param[out] checkpoint_state + * Checkpoint states that contains the states of the training session. + * + *

\snippet{doc} snippets.dox OrtStatus Return Value + * + *

ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ + * OrtCheckpointState** checkpoint_state); + */ + private static native long loadCheckpoint(long apiHandle, long trainingApiHandle, String path) + throws OrtException; + + private native void close(long trainingApiHandle, long nativeHandle); + } +} diff --git a/java/src/main/java/ai/onnxruntime/providers/package-info.java b/java/src/main/java/ai/onnxruntime/providers/package-info.java new file mode 100644 index 0000000000..1f1e70a589 --- /dev/null +++ b/java/src/main/java/ai/onnxruntime/providers/package-info.java @@ -0,0 +1,7 @@ +/* + * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ + +/** Classes for controlling the behaviour of ONNX Runtime Execution Providers. */ +package ai.onnxruntime.providers; diff --git a/java/src/main/native/ai_onnxruntime_OnnxRuntime.c b/java/src/main/native/ai_onnxruntime_OnnxRuntime.c index a209f2df2b..d0c4be87d9 100644 --- a/java/src/main/native/ai_onnxruntime_OnnxRuntime.c +++ b/java/src/main/native/ai_onnxruntime_OnnxRuntime.c @@ -17,7 +17,19 @@ JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OnnxRuntime_initialiseAPIBase(JNIEnv jint apiVersion) { (void)jniEnv; (void)clazz; // required JNI parameters not needed by functions which don't call back into Java. const OrtApi* ortPtr = OrtGetApiBase()->GetApi((uint32_t)apiVersion); - return (jlong)ortPtr; + return (jlong) ortPtr; +} +/* + * Class: ai_onnxruntime_OnnxRuntime + * Method: initialiseTrainingAPIBase + * Signature: (JI)J + */ +JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OnnxRuntime_initialiseTrainingAPIBase + (JNIEnv * jniEnv, jclass clazz, jlong apiHandle, jint trainingApiVersion) { + (void)jniEnv; (void)clazz; // required JNI parameters not needed by functions which don't call back into Java. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainingApi = api->GetTrainingApi((uint32_t)trainingApiVersion); + return (jlong) trainingApi; } /* diff --git a/java/src/main/native/ai_onnxruntime_OrtSession_RunOptions.c b/java/src/main/native/ai_onnxruntime_OrtSession_RunOptions.c index 9086082e69..0ab9bbb889 100644 --- a/java/src/main/native/ai_onnxruntime_OrtSession_RunOptions.c +++ b/java/src/main/native/ai_onnxruntime_OrtSession_RunOptions.c @@ -108,6 +108,22 @@ JNIEXPORT jstring JNICALL Java_ai_onnxruntime_OrtSession_00024RunOptions_getRunT } } +/* + * Class: ai_onnxruntime_OrtSession_RunOptions + * Method: addRunConfigEntry + * Signature: (JJLjava/lang/String;Ljava/lang/String;)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_00024RunOptions_addRunConfigEntry + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong nativeHandle, jstring keyStr, jstring valueStr) { + (void) jobj; // Required JNI parameters not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*) apiHandle; + const char* key = (*jniEnv)->GetStringUTFChars(jniEnv, keyStr, NULL); + const char* value = (*jniEnv)->GetStringUTFChars(jniEnv, valueStr, NULL); + checkOrtStatus(jniEnv, api, api->AddRunConfigEntry((OrtRunOptions*) nativeHandle, key, value)); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, keyStr, key); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, valueStr, value); +} + /* * Class: ai_onnxruntime_OrtSession_RunOptions * Method: setTerminate diff --git a/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c b/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c new file mode 100644 index 0000000000..591dae7888 --- /dev/null +++ b/java/src/main/native/ai_onnxruntime_OrtTrainingSession.c @@ -0,0 +1,762 @@ +/* + * Copyright (c) 2022 Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +#include +#include +#include +#include "OrtJniUtil.h" +#include "onnxruntime/core/session/onnxruntime_c_api.h" +#include "onnxruntime_training_c_api.h" +#include "ai_onnxruntime_OrtTrainingSession.h" + +#ifdef _WIN32 +wchar_t* copyAndPad(JNIEnv * jniEnv, jstring javaStr) { + // The output of GetStringChars is not null-terminated, so we copy it and add a terminator + const jchar* charArr = (*jniEnv)->GetStringChars(jniEnv, javaStr, NULL); + size_t strLength = (*jniEnv)->GetStringLength(jniEnv, javaStr); + wchar_t* outputStr = (wchar_t*)calloc(strLength + 1, sizeof(wchar_t)); + if (outputStr != NULL) { + wcsncpy_s(outputStr, strLength + 1, (const wchar_t*)charArr, strLength); + } else { + throwOrtException(jniEnv, 1, "Not enough memory"); + } + (*jniEnv)->ReleaseStringChars(jniEnv, javaStr, charArr); + return outputStr; +} +#endif + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: createTrainingSession + * Signature: (JJJJJLjava/lang/String;Ljava/lang/String;Ljava/lang/String;)J + */ +JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OrtTrainingSession_createTrainingSession + (JNIEnv * jniEnv, jclass clazz, jlong apiHandle, jlong trainApiHandle, + jlong envHandle, jlong optionsHandle, jlong checkpointHandle, + jstring trainPath, jstring evalPath, jstring optimizerPath) { + (void) clazz; // Required JNI parameters not needed by functions which don't need to access their host class. + + // evalPath and optimizerPath could be NULL, as that is used to signal that those models + // should not be loaded, which induces some juggling to avoid calling JNI methods with a NULL + // pointer. trainPath cannot be null, as in that case a Java exception is thrown before this + // method is called. + + const OrtApi* api = (const OrtApi*) apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*) trainApiHandle; + const OrtEnv* env = (const OrtEnv*) envHandle; + const OrtSessionOptions* options = (const OrtSessionOptions*) optionsHandle; + OrtCheckpointState* checkpoint = (OrtCheckpointState*) checkpointHandle; + + OrtTrainingSession* session = NULL; + +#ifdef _WIN32 + // The output of GetStringChars is not null-terminated, so we copy it and add a terminator + wchar_t* trainStr = copyAndPad(jniEnv, trainPath); + if (trainStr == NULL) { + // nothing to cleanup, return zero as exception has been thrown in Java + return 0L; + } + wchar_t* evalStr = NULL; + if (evalPath != NULL) { + evalStr = copyAndPad(jniEnv, evalPath); + if (evalStr == NULL) { + // exception has been thrown in Java, go to cleanup and return null. + goto cleanupTrain; + } + } + wchar_t* optimizerStr = NULL; + if (optimizerPath == NULL) { + optimizerStr = copyAndPad(jniEnv, optimizerPath); + if (optimizerStr == NULL) { + // exception has been thrown in Java, go to cleanup and return null. + goto cleanupEval; + } + } + checkOrtStatus(jniEnv, api, trainApi->CreateTrainingSession(env, options, checkpoint, trainStr, evalStr, optimizerStr, &session)); + if (optimizerStr != NULL) { + free(optimizerStr); + } +cleanupEval: + if (evalStr != NULL) { + free(evalStr); + } +cleanupTrain: + free(trainStr); +#else + // GetStringUTFChars is null terminated, so can be used directly + const char* trainStr = (*jniEnv)->GetStringUTFChars(jniEnv, trainPath, NULL); + const char* evalStr = evalPath == NULL ? NULL : (*jniEnv)->GetStringUTFChars(jniEnv, evalPath, NULL); + const char* optimizerStr = optimizerPath == NULL ? NULL : (*jniEnv)->GetStringUTFChars(jniEnv, optimizerPath, NULL); + checkOrtStatus(jniEnv, api, trainApi->CreateTrainingSession(env, options, checkpoint, trainStr, evalStr, optimizerStr, &session)); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, trainPath, trainStr); + if (evalPath != NULL) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, evalPath, evalStr); + } + if (optimizerPath != NULL) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, optimizerPath, optimizerStr); + } +#endif + + return (jlong) session; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: closeSession + * Signature: (JJ)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_closeSession + (JNIEnv * jniEnv, jobject jobj, jlong trainHandle, jlong nativeHandle) { + (void)jniEnv; (void)jobj; // Required JNI parameters not needed by functions which don't need to access their host object. + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainHandle; + trainApi->ReleaseTrainingSession((OrtTrainingSession*)nativeHandle); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: saveCheckpoint + * Signature: (JJJLjava/lang/String;Z)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_saveCheckpoint + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainingApiHandle, jlong nativeHandle, jstring outputPath, jboolean overwrite) { + (void) jobj; // Required JNI parameters not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*) apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*) trainingApiHandle; + + const OrtTrainingSession* trainSession = (const OrtTrainingSession*) nativeHandle; + +#ifdef _WIN32 + // The output of GetStringChars is not null-terminated, so we copy it and add a terminator + const jchar* cPath = (*jniEnv)->GetStringChars(jniEnv, outputPath, NULL); + size_t stringLength = (*jniEnv)->GetStringLength(jniEnv, outputPath); + wchar_t* newString = (wchar_t*)calloc(stringLength + 1, sizeof(wchar_t)); + if (newString == NULL) { + (*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath); + throwOrtException(jniEnv, 1, "Not enough memory"); + } else { + wcsncpy_s(newString, stringLength + 1, (const wchar_t*)cPath, stringLength); + checkOrtStatus(jniEnv, api, + trainApi->SaveCheckpoint(newString, trainSession, overwrite)); + free(newString); + (*jniEnv)->ReleaseStringChars(jniEnv, outputPath, cPath); + } +#else + // GetStringUTFChars is null terminated, so can be used directly + const char* cPath = (*jniEnv)->GetStringUTFChars(jniEnv, outputPath, NULL); + checkOrtStatus(jniEnv, api, trainApi->SaveCheckpoint(cPath, trainSession, overwrite)); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, outputPath, cPath); +#endif +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: getTrainInputNames + * Signature: (JJJJ)[Ljava/lang/String; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_getTrainInputNames + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong sessionHandle, jlong allocatorHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + const OrtTrainingSession* trainSession = (const OrtTrainingSession*)sessionHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + + // Setup + jclass stringClazz = (*jniEnv)->FindClass(jniEnv, "java/lang/String"); + + // Get the number of inputs + size_t numInputs = 0; + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetTrainingModelInputCount(trainSession, &numInputs)); + if (code != ORT_OK) { + return NULL; + } + + int32_t numInputsInt = (int32_t) numInputs; + if (numInputs != (size_t) numInputsInt) { + throwOrtException(jniEnv, 1, "Too many inputs, expected less than 2^31"); + } + + // Allocate the return array + jobjectArray array = (*jniEnv)->NewObjectArray(jniEnv, numInputsInt, stringClazz, NULL); + for (int32_t i = 0; i < numInputsInt; i++) { + // Read out the input name and convert it to a java.lang.String + char* inputName = NULL; + code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetTrainingModelInputName(trainSession, i, allocator, &inputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + jstring name = (*jniEnv)->NewStringUTF(jniEnv, inputName); + (*jniEnv)->SetObjectArrayElement(jniEnv, array, i, name); + code = checkOrtStatus(jniEnv, api, api->AllocatorFree(allocator, inputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + } + + return array; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: getTrainOutputNames + * Signature: (JJJJ)[Ljava/lang/String; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_getTrainOutputNames + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong sessionHandle, jlong allocatorHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + const OrtTrainingSession* trainSession = (const OrtTrainingSession*)sessionHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + + // Setup + jclass stringClazz = (*jniEnv)->FindClass(jniEnv, "java/lang/String"); + + // Get the number of outputs + size_t numOutputs = 0; + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetTrainingModelOutputCount(trainSession, &numOutputs)); + if (code != ORT_OK) { + return NULL; + } + + int32_t numOutputsInt = (int32_t) numOutputs; + if (numOutputs != (size_t) numOutputsInt) { + throwOrtException(jniEnv, 1, "Too many outputs, expected less than 2^31"); + } + + // Allocate the return array + jobjectArray array = (*jniEnv)->NewObjectArray(jniEnv, numOutputsInt, stringClazz, NULL); + for (int32_t i = 0; i < numOutputsInt; i++) { + // Read out the output name and convert it to a java.lang.String + char* outputName = NULL; + code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetTrainingModelOutputName(trainSession, i, allocator, &outputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + jstring name = (*jniEnv)->NewStringUTF(jniEnv, outputName); + (*jniEnv)->SetObjectArrayElement(jniEnv, array, i, name); + code = checkOrtStatus(jniEnv, api, api->AllocatorFree(allocator, outputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + } + + return array; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: getEvalInputNames + * Signature: (JJJJ)[Ljava/lang/String; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_getEvalInputNames + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong sessionHandle, jlong allocatorHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + const OrtTrainingSession* trainSession = (const OrtTrainingSession*)sessionHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + + // Setup + jclass stringClazz = (*jniEnv)->FindClass(jniEnv, "java/lang/String"); + + // Get the number of inputs + size_t numInputs = 0; + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetEvalModelInputCount(trainSession, &numInputs)); + if (code != ORT_OK) { + return NULL; + } + + int32_t numInputsInt = (int32_t) numInputs; + if (numInputs != (size_t) numInputsInt) { + throwOrtException(jniEnv, 1, "Too many inputs, expected less than 2^31"); + } + + // Allocate the return array + jobjectArray array = (*jniEnv)->NewObjectArray(jniEnv, numInputsInt, stringClazz, NULL); + for (int32_t i = 0; i < numInputsInt; i++) { + // Read out the input name and convert it to a java.lang.String + char* inputName = NULL; + code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetEvalModelInputName(trainSession, i, allocator, &inputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + jstring name = (*jniEnv)->NewStringUTF(jniEnv, inputName); + (*jniEnv)->SetObjectArrayElement(jniEnv, array, i, name); + code = checkOrtStatus(jniEnv, api, api->AllocatorFree(allocator, inputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + } + + return array; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: getEvalOutputNames + * Signature: (JJJJ)[Ljava/lang/String; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_getEvalOutputNames + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong sessionHandle, jlong allocatorHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + const OrtTrainingSession* trainSession = (const OrtTrainingSession*)sessionHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + + // Setup + jclass stringClazz = (*jniEnv)->FindClass(jniEnv, "java/lang/String"); + + // Get the number of outputs + size_t numOutputs = 0; + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetEvalModelOutputCount(trainSession, &numOutputs)); + if (code != ORT_OK) { + return NULL; + } + + int32_t numOutputsInt = (int32_t) numOutputs; + if (numOutputs != (size_t) numOutputsInt) { + throwOrtException(jniEnv, 1, "Too many outputs, expected less than 2^31"); + } + + // Allocate the return array + jobjectArray array = (*jniEnv)->NewObjectArray(jniEnv, numOutputsInt, stringClazz, NULL); + for (int32_t i = 0; i < numOutputsInt; i++) { + // Read out the output name and convert it to a java.lang.String + char* outputName = NULL; + code = checkOrtStatus(jniEnv, api, trainApi->TrainingSessionGetEvalModelOutputName(trainSession, i, allocator, &outputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + jstring name = (*jniEnv)->NewStringUTF(jniEnv, outputName); + (*jniEnv)->SetObjectArrayElement(jniEnv, array, i, name); + code = checkOrtStatus(jniEnv, api, api->AllocatorFree(allocator, outputName)); + if (code != ORT_OK) { + // break out on error, return array and let Java throw the exception. + break; + } + } + + return array; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: lazyResetGrad + * Signature: (JJJ)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_lazyResetGrad + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + checkOrtStatus(jniEnv, api, trainApi->LazyResetGrad(trainSession)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: trainStep + * Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;JJ)[Lai/onnxruntime/OnnxValue; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, + jlong nativeHandle, jlong allocatorHandle, jobjectArray inputNamesArr, jlongArray inputHandles, jlong numInputs, + jobjectArray outputNamesArr, jlong numOutputs, jlong runOptionsHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + OrtRunOptions* runOptions = (OrtRunOptions*)runOptionsHandle; + + jobjectArray outputArray = NULL; + + // Create the buffers for the Java input & output strings, and the input pointers + const char** inputNames = malloc(sizeof(char*) * numInputs); + if (inputNames == NULL) { + // Nothing to cleanup, return and throw exception + return outputArray; + } + const char** outputNames = malloc(sizeof(char*) * numOutputs); + if (outputNames == NULL) { + goto cleanup_input_names; + } + jobject* javaInputStrings = malloc(sizeof(jobject) * numInputs); + if (javaInputStrings == NULL) { + goto cleanup_output_names; + } + jobject* javaOutputStrings = malloc(sizeof(jobject) * numOutputs); + if (javaOutputStrings == NULL) { + goto cleanup_java_input_strings; + } + const OrtValue** inputValuePtrs = malloc(sizeof(OrtValue*) * numInputs); + if (inputValuePtrs == NULL) { + goto cleanup_java_output_strings; + } + OrtValue** outputValues = malloc(sizeof(OrtValue*) * numOutputs); + if (outputValues == NULL) { + goto cleanup_input_values; + } + + // Extract a C array of longs which are pointers to the input tensors. + // The Java-side objects store native pointers as 64-bit longs, and on 32-bit systems + // we cannot cast the long array to a pointer array as they are different sizes, + // so we copy the longs applying the appropriate cast. + jlong* inputValueLongs = (*jniEnv)->GetLongArrayElements(jniEnv, inputHandles, NULL); + + // Extract the names and native pointers of the input values. + for (int i = 0; i < numInputs; i++) { + javaInputStrings[i] = (*jniEnv)->GetObjectArrayElement(jniEnv, inputNamesArr, i); + inputNames[i] = (*jniEnv)->GetStringUTFChars(jniEnv, javaInputStrings[i], NULL); + inputValuePtrs[i] = (OrtValue*)inputValueLongs[i]; + } + + // Release the java array copy of pointers to the tensors. + (*jniEnv)->ReleaseLongArrayElements(jniEnv, inputHandles, inputValueLongs, JNI_ABORT); + + // Extract the names of the output values. + for (int i = 0; i < numOutputs; i++) { + javaOutputStrings[i] = (*jniEnv)->GetObjectArrayElement(jniEnv, outputNamesArr, i); + outputNames[i] = (*jniEnv)->GetStringUTFChars(jniEnv, javaOutputStrings[i], NULL); + outputValues[i] = NULL; + } + + // Actually score the inputs. + //ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + // size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + // size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->TrainStep(trainSession, runOptions, + numInputs, (const OrtValue* const*)inputValuePtrs, + numOutputs, outputValues)); + if (code != ORT_OK) { + goto cleanup_output_values; + } + + // Construct the output array of ONNXValues + jclass onnxValueClass = (*jniEnv)->FindClass(jniEnv, "ai/onnxruntime/OnnxValue"); + outputArray = (*jniEnv)->NewObjectArray(jniEnv, safecast_int64_to_jsize(numOutputs), onnxValueClass, NULL); + + // Convert the output tensors into ONNXValues + for (int i = 0; i < numOutputs; i++) { + if (outputValues[i] != NULL) { + jobject onnxValue = convertOrtValueToONNXValue(jniEnv, api, allocator, outputValues[i]); + if (onnxValue == NULL) { + break; // go to cleanup, exception thrown + } + (*jniEnv)->SetObjectArrayElement(jniEnv, outputArray, i, onnxValue); + } + } + + // Note these gotos are in a specific order so they mirror the allocation pattern above. + // They must be changed if the allocation code is rearranged. + cleanup_output_values: + free(outputValues); + + // Release the Java output strings + for (int i = 0; i < numOutputs; i++) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, javaOutputStrings[i], outputNames[i]); + } + + // Release the Java input strings + for (int i = 0; i < numInputs; i++) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, javaInputStrings[i], inputNames[i]); + } + + // Release the buffers + cleanup_input_values: + free((void*)inputValuePtrs); + cleanup_java_output_strings: + free(javaOutputStrings); + cleanup_java_input_strings: + free(javaInputStrings); + cleanup_output_names: + free((void*)outputNames); + cleanup_input_names: + free((void*)inputNames); + + return outputArray; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: evalStep + * Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;JJ)[Lai/onnxruntime/OnnxValue; + */ +JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, + jlong nativeHandle, jlong allocatorHandle, jobjectArray inputNamesArr, jlongArray inputHandles, jlong numInputs, + jobjectArray outputNamesArr, jlong numOutputs, jlong runOptionsHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtAllocator* allocator = (OrtAllocator*)allocatorHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + OrtRunOptions* runOptions = (OrtRunOptions*)runOptionsHandle; + + jobjectArray outputArray = NULL; + + // Create the buffers for the Java input & output strings, and the input pointers + const char** inputNames = malloc(sizeof(char*) * numInputs); + if (inputNames == NULL) { + // Nothing to cleanup, return and throw exception + return outputArray; + } + const char** outputNames = malloc(sizeof(char*) * numOutputs); + if (outputNames == NULL) { + goto cleanup_input_names; + } + jobject* javaInputStrings = malloc(sizeof(jobject) * numInputs); + if (javaInputStrings == NULL) { + goto cleanup_output_names; + } + jobject* javaOutputStrings = malloc(sizeof(jobject) * numOutputs); + if (javaOutputStrings == NULL) { + goto cleanup_java_input_strings; + } + const OrtValue** inputValuePtrs = malloc(sizeof(OrtValue*) * numInputs); + if (inputValuePtrs == NULL) { + goto cleanup_java_output_strings; + } + OrtValue** outputValues = malloc(sizeof(OrtValue*) * numOutputs); + if (outputValues == NULL) { + goto cleanup_input_values; + } + + // Extract a C array of longs which are pointers to the input tensors. + // The Java-side objects store native pointers as 64-bit longs, and on 32-bit systems + // we cannot cast the long array to a pointer array as they are different sizes, + // so we copy the longs applying the appropriate cast. + jlong* inputValueLongs = (*jniEnv)->GetLongArrayElements(jniEnv, inputHandles, NULL); + + // Extract the names and native pointers of the input values. + for (int i = 0; i < numInputs; i++) { + javaInputStrings[i] = (*jniEnv)->GetObjectArrayElement(jniEnv, inputNamesArr, i); + inputNames[i] = (*jniEnv)->GetStringUTFChars(jniEnv, javaInputStrings[i], NULL); + inputValuePtrs[i] = (OrtValue*)inputValueLongs[i]; + } + + // Release the java array copy of pointers to the tensors. + (*jniEnv)->ReleaseLongArrayElements(jniEnv, inputHandles, inputValueLongs, JNI_ABORT); + + // Extract the names of the output values. + for (int i = 0; i < numOutputs; i++) { + javaOutputStrings[i] = (*jniEnv)->GetObjectArrayElement(jniEnv, outputNamesArr, i); + outputNames[i] = (*jniEnv)->GetStringUTFChars(jniEnv, javaOutputStrings[i], NULL); + outputValues[i] = NULL; + } + + // Actually score the inputs. + //ORT_API2_STATUS(EvalStep, _In_ const OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + // size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + // size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + OrtErrorCode code = checkOrtStatus(jniEnv, api, trainApi->EvalStep(trainSession, runOptions, + numInputs, (const OrtValue* const*)inputValuePtrs, + numOutputs, outputValues)); + if (code != ORT_OK) { + goto cleanup_output_values; + } + + // Construct the output array of ONNXValues + jclass onnxValueClass = (*jniEnv)->FindClass(jniEnv, "ai/onnxruntime/OnnxValue"); + outputArray = (*jniEnv)->NewObjectArray(jniEnv, safecast_int64_to_jsize(numOutputs), onnxValueClass, NULL); + + // Convert the output tensors into ONNXValues + for (int i = 0; i < numOutputs; i++) { + if (outputValues[i] != NULL) { + jobject onnxValue = convertOrtValueToONNXValue(jniEnv, api, allocator, outputValues[i]); + if (onnxValue == NULL) { + break; // go to cleanup, exception thrown + } + (*jniEnv)->SetObjectArrayElement(jniEnv, outputArray, i, onnxValue); + } + } + + // Note these gotos are in a specific order so they mirror the allocation pattern above. + // They must be changed if the allocation code is rearranged. + cleanup_output_values: + free(outputValues); + + // Release the Java output strings + for (int i = 0; i < numOutputs; i++) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, javaOutputStrings[i], outputNames[i]); + } + + // Release the Java input strings + for (int i = 0; i < numInputs; i++) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, javaInputStrings[i], inputNames[i]); + } + + // Release the buffers + cleanup_input_values: + free((void*)inputValuePtrs); + cleanup_java_output_strings: + free(javaOutputStrings); + cleanup_java_input_strings: + free(javaInputStrings); + cleanup_output_names: + free((void*)outputNames); + cleanup_input_names: + free((void*)inputNames); + + return outputArray; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: setSeed + * Signature: (JJJF)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_setSeed + (JNIEnv * jniEnv, jclass clazz, jlong apiHandle, jlong trainApiHandle, jlong seed) { + (void)clazz; // Required JNI parameter not needed by functions which don't need to access their host class. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + checkOrtStatus(jniEnv, api, trainApi->SetSeed(seed)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: setLearningRate + * Signature: (JJJF)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_setLearningRate + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle, jfloat learningRate) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + checkOrtStatus(jniEnv, api, trainApi->SetLearningRate(trainSession, learningRate)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: getLearningRate + * Signature: (JJJ)F + */ +JNIEXPORT jfloat JNICALL Java_ai_onnxruntime_OrtTrainingSession_getLearningRate + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + jfloat learningRate = 0.0f; + checkOrtStatus(jniEnv, api, trainApi->GetLearningRate(trainSession, &learningRate)); + return learningRate; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: optimizerStep + * Signature: (JJJJ)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_optimizerStep + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle, jlong runOptionsHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + const OrtRunOptions* options = (const OrtRunOptions*) runOptionsHandle; + checkOrtStatus(jniEnv, api, trainApi->OptimizerStep(trainSession, options)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: registerLinearLRScheduler + * Signature: (JJJJJF)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_registerLinearLRScheduler + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle, jlong warmupSteps, jlong totalSteps, jfloat initialLearningRate) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + checkOrtStatus(jniEnv, api, trainApi->RegisterLinearLRScheduler(trainSession, warmupSteps, totalSteps, initialLearningRate)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: schedulerStep + * Signature: (JJJ)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_schedulerStep + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + checkOrtStatus(jniEnv, api, trainApi->SchedulerStep(trainSession)); +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession + * Method: exportModelForInference + * Signature: (JJJJLjava/lang/String;[Ljava/lang/String;)V + */ +#ifdef _MSC_VER +#pragma warning(push) +// C4090: 'operation' : different 'modifier' qualifiers +// Freeing 'outputNames' erroneously triggers this warning, it is fixed in VC 2022 and can be removed when that is the baseline compiler. +#pragma warning(disable : 4090) +#endif // _MSC_VER +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_exportModelForInference + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong trainApiHandle, jlong nativeHandle, jstring outputPath, jlong numOutputs, jobjectArray outputNamesArr) { + (void)jobj; // Required JNI parameter not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*)apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*)trainApiHandle; + OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle; + + // prep output names array + const char** outputNames = malloc(sizeof(char*) * numOutputs); + if (outputNames == NULL) { + throwOrtException(jniEnv, 1, "Not enough memory"); + return; + } + jobject* javaOutputStrings = malloc(sizeof(jobject) * numOutputs); + if (javaOutputStrings == NULL) { + throwOrtException(jniEnv, 1, "Not enough memory"); + free(outputNames); + return; + } + // Extract the names of the output values. + for (int i = 0; i < numOutputs; i++) { + javaOutputStrings[i] = (*jniEnv)->GetObjectArrayElement(jniEnv, outputNamesArr, i); + outputNames[i] = (*jniEnv)->GetStringUTFChars(jniEnv, javaOutputStrings[i], NULL); + } + +#ifdef _WIN32 + // The output of GetStringChars is not null-terminated, so we copy it and add a terminator + wchar_t* outputStr = copyAndPad(jniEnv, outputPath); + if (outputStr == NULL) { + goto cleanup_array; + } + checkOrtStatus(jniEnv, api, trainApi->ExportModelForInferencing(trainSession, outputStr, numOutputs, outputNames)); + free(outputStr); +#else + // GetStringUTFChars is null terminated, so can be used directly + const char* outputStr = (*jniEnv)->GetStringUTFChars(jniEnv, outputPath, NULL); + checkOrtStatus(jniEnv, api, trainApi->ExportModelForInferencing(trainSession, outputStr, numOutputs, outputNames)); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, outputPath, outputStr); + goto cleanup_array; // Only used in the WIN32 branch, but gcc complains we don't use this label otherwise +#endif + +cleanup_array: + // Release the Java output strings + for (int i = 0; i < numOutputs; i++) { + (*jniEnv)->ReleaseStringUTFChars(jniEnv, javaOutputStrings[i], outputNames[i]); + } + free(javaOutputStrings); + free(outputNames); +} +#ifdef _MSC_VER +#pragma warning(pop) +#endif // _MSC_VER diff --git a/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c b/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c new file mode 100644 index 0000000000..7678f28ffa --- /dev/null +++ b/java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +#include +#include +#include "onnxruntime/core/session/onnxruntime_c_api.h" +#include "onnxruntime_training_c_api.h" +#include "OrtJniUtil.h" +#include "ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.h" + +/* + * Class: ai_onnxruntime_OrtTrainingSession_OrtCheckpointState + * Method: loadCheckpoint + * Signature: (JJLjava/lang/String;)J + */ +JNIEXPORT jlong JNICALL Java_ai_onnxruntime_OrtTrainingSession_00024OrtCheckpointState_loadCheckpoint + (JNIEnv * jniEnv, jclass jclazz, jlong apiHandle, jlong trainingApiHandle, jstring directory) { + (void) jclazz; // Required JNI parameters not needed by functions which don't need to access their host object. + const OrtApi* api = (const OrtApi*) apiHandle; + const OrtTrainingApi* trainApi = (const OrtTrainingApi*) trainingApiHandle; + + OrtCheckpointState* checkpoint = NULL; + +#ifdef _WIN32 + const jchar* cPath = (*jniEnv)->GetStringChars(jniEnv, directory, NULL); + size_t stringLength = (*jniEnv)->GetStringLength(jniEnv, directory); + wchar_t* newString = (wchar_t*)calloc(stringLength + 1, sizeof(wchar_t)); + if (newString == NULL) { + (*jniEnv)->ReleaseStringChars(jniEnv, directory, cPath); + throwOrtException(jniEnv, 1, "Not enough memory"); + return 0; + } + wcsncpy_s(newString, stringLength + 1, (const wchar_t*)cPath, stringLength); + checkOrtStatus(jniEnv, api, + trainApi->LoadCheckpoint(newString, &checkpoint)); + free(newString); + (*jniEnv)->ReleaseStringChars(jniEnv, directory, cPath); +#else + const char* cPath = (*jniEnv)->GetStringUTFChars(jniEnv, directory, NULL); + checkOrtStatus(jniEnv, api, trainApi->LoadCheckpoint(cPath, &checkpoint)); + (*jniEnv)->ReleaseStringUTFChars(jniEnv, directory, cPath); +#endif + + return (jlong) checkpoint; +} + +/* + * Class: ai_onnxruntime_OrtTrainingSession_OrtCheckpointState + * Method: close + * Signature: (JJ)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_00024OrtCheckpointState_close + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong handle) { + (void) jniEnv; (void) jobj; // Required JNI parameters not needed by functions which don't need to access their host object. + const OrtTrainingApi* api = (const OrtTrainingApi*) apiHandle; + api->ReleaseCheckpointState((OrtCheckpointState*) handle); +} diff --git a/java/src/test/java/ai/onnxruntime/InferenceTest.java b/java/src/test/java/ai/onnxruntime/InferenceTest.java index d545292aea..8fc5fe2daf 100644 --- a/java/src/test/java/ai/onnxruntime/InferenceTest.java +++ b/java/src/test/java/ai/onnxruntime/InferenceTest.java @@ -59,10 +59,6 @@ public class InferenceTest { private static final OrtEnvironment env = OrtEnvironment.getEnvironment(); - public static Path getResourcePath(String path) { - return new File(InferenceTest.class.getResource(path).getFile()).toPath(); - } - @Test public void environmentTest() { // Checks that the environment instance is the same. diff --git a/java/src/test/java/ai/onnxruntime/SparseTensorTest.java b/java/src/test/java/ai/onnxruntime/SparseTensorTest.java index 7a3abb7bf7..9b0533f5c4 100644 --- a/java/src/test/java/ai/onnxruntime/SparseTensorTest.java +++ b/java/src/test/java/ai/onnxruntime/SparseTensorTest.java @@ -4,7 +4,6 @@ */ package ai.onnxruntime; -import static ai.onnxruntime.InferenceTest.getResourcePath; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -25,7 +24,8 @@ public class SparseTensorTest { @Test public void testCSRC() throws OrtException { - String modelPath = getResourcePath("/generic_sparse_to_dense_matmul.onnx").toString(); + String modelPath = + TestHelpers.getResourcePath("/generic_sparse_to_dense_matmul.onnx").toString(); try (OrtEnvironment env = OrtEnvironment.getEnvironment(); OrtSession.SessionOptions options = new OrtSession.SessionOptions()) { try (OrtSession session = env.createSession(modelPath, options)) { @@ -205,7 +205,8 @@ public class SparseTensorTest { @Test public void testCOO() throws OrtException { - String modelPath = getResourcePath("/generic_sparse_to_dense_matmul.onnx").toString(); + String modelPath = + TestHelpers.getResourcePath("/generic_sparse_to_dense_matmul.onnx").toString(); try (OrtEnvironment env = OrtEnvironment.getEnvironment(); OrtSession.SessionOptions options = new OrtSession.SessionOptions()) { try (OrtSession session = env.createSession(modelPath, options)) { @@ -391,7 +392,7 @@ public class SparseTensorTest { @Test public void testCOOOutput() throws OrtException { - String modelPath = getResourcePath("/sparse_initializer_as_output.onnx").toString(); + String modelPath = TestHelpers.getResourcePath("/sparse_initializer_as_output.onnx").toString(); try (OrtEnvironment env = OrtEnvironment.getEnvironment(); OrtSession.SessionOptions options = new OrtSession.SessionOptions()) { try (OrtSession session = env.createSession(modelPath, options)) { diff --git a/java/src/test/java/ai/onnxruntime/TestHelpers.java b/java/src/test/java/ai/onnxruntime/TestHelpers.java index 8af12cf4dd..b82cb07aee 100644 --- a/java/src/test/java/ai/onnxruntime/TestHelpers.java +++ b/java/src/test/java/ai/onnxruntime/TestHelpers.java @@ -13,9 +13,11 @@ import java.io.IOException; import java.io.InputStream; import java.io.UncheckedIOException; import java.nio.ByteBuffer; +import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.regex.Pattern; @@ -26,6 +28,10 @@ public class TestHelpers { private static final Pattern LOAD_PATTERN = Pattern.compile("[,\\[\\] ]"); + static void deleteDirectoryTree(Path input) throws IOException { + Files.walk(input).sorted(Comparator.reverseOrder()).map(Path::toFile).forEach(File::delete); + } + static boolean[] toPrimitiveBoolean(List input) { boolean[] output = new boolean[input.size()]; @@ -251,7 +257,7 @@ public class TestHelpers { } public static Path getResourcePath(String path) { - return new File(InferenceTest.class.getResource(path).getFile()).toPath(); + return new File(TestHelpers.class.getResource(path).getFile()).toPath(); } public static float[] loadTensorFromFile(Path filename) { diff --git a/java/src/test/java/ai/onnxruntime/TrainingTest.java b/java/src/test/java/ai/onnxruntime/TrainingTest.java new file mode 100644 index 0000000000..e5140d7110 --- /dev/null +++ b/java/src/test/java/ai/onnxruntime/TrainingTest.java @@ -0,0 +1,256 @@ +/* + * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +package ai.onnxruntime; + +import ai.onnxruntime.OrtTrainingSession.OrtCheckpointState; +import ai.onnxruntime.TensorInfo.OnnxTensorType; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.FloatBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfSystemProperty; + +/** Tests for the ORT training apis. */ +@EnabledIfSystemProperty(named = "ENABLE_TRAINING", matches = "1") +public class TrainingTest { + + private static final OrtEnvironment env = OrtEnvironment.getEnvironment(); + + @Test + public void testLoadCheckpoint() throws OrtException { + Path ckptPath = TestHelpers.getResourcePath("/checkpoint.ckpt"); + try (OrtCheckpointState ckpt = OrtCheckpointState.loadCheckpoint(ckptPath)) { + // Must be non-null, exists so the try block isn't empty as this call will + // throw if it fails, and throwing errors the test + Assertions.assertNotNull(ckpt); + } + } + + @Test + public void testCreateTrainingSession() throws OrtException { + String ckptPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + try (OrtTrainingSession trainingSession = + env.createTrainingSession(ckptPath, trainPath, null, null)) { + Assertions.assertNotNull(trainingSession); + Set inputNames = trainingSession.getTrainInputNames(); + Assertions.assertFalse(inputNames.isEmpty()); + Set outputNames = trainingSession.getTrainOutputNames(); + Assertions.assertFalse(outputNames.isEmpty()); + } + } + + // this test is not enabled as ORT Java doesn't support supplying an output buffer + @Disabled + @Test + public void TestTrainingSessionTrainStep() throws OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + float[] expectedOutput = + TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/loss_1.out")); + float[] input = TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/input-0.in")); + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, null)) { + int[] labels = {1, 1}; + + // Run train step with pinned inputs and pinned outputs + Map pinnedInputs = new HashMap<>(); + Map outputMap = new HashMap<>(); + try { + // Create inputs + long[] inputShape = {2, 784}; + pinnedInputs.put( + "input-0", OnnxTensor.createTensor(env, OrtUtil.reshape(input, inputShape))); + + // long[] labelsShape = {2}; + pinnedInputs.put("labels", OnnxTensor.createTensor(env, labels)); + + // Prepare output buffer + FloatBuffer output = + ByteBuffer.allocateDirect(4 * expectedOutput.length) + .order(ByteOrder.nativeOrder()) + .asFloatBuffer(); + OnnxTensor outputTensor = + OnnxTensor.createTensor(env, output, new long[expectedOutput.length]); + outputMap.put("onnx::loss::21273", outputTensor); + /* Disabled as we haven't implemented this yet + try (trainingSession.trainStep(pinnedInputs, outputMap)) { + Assertions.assertArrayEquals(expectedOutput, (float[]) outputTensor.getValue(), 1e-3f); + } + */ + } finally { + OnnxValue.close(outputMap); + OnnxValue.close(pinnedInputs); + } + } + } + + void runTrainStep(OrtTrainingSession trainingSession) throws OrtException { + float[] expectedOutput = + TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/loss_1.out")); + float[] input = TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/input-0.in")); + int[] labels = {1, 1}; + + // Run inference with pinned inputs and pinned outputs + + // Create inputs + Map pinnedInputs = new HashMap<>(); + try { + long[] inputShape = {2, 784}; + pinnedInputs.put("input-0", OnnxTensor.createTensor(env, OrtUtil.reshape(input, inputShape))); + + // long[] labelsShape = {2}; + pinnedInputs.put("labels", OnnxTensor.createTensor(env, labels)); + + try (OrtSession.Result firstOutput = trainingSession.trainStep(pinnedInputs)) { + Assertions.assertTrue(firstOutput.size() > 0); + } + trainingSession.lazyResetGrad(); + try (OrtSession.Result secondOutputs = trainingSession.trainStep(pinnedInputs)) { + OnnxValue outputBuffer = secondOutputs.get(0); + + Assertions.assertEquals(secondOutputs.get("onnx::loss::21273").get(), outputBuffer); + Assertions.assertTrue(outputBuffer instanceof OnnxTensor); + + OnnxTensor outLabelTensor = (OnnxTensor) outputBuffer; + Assertions.assertEquals( + OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, outLabelTensor.getInfo().onnxType); + Assertions.assertNotNull(outLabelTensor); + Assertions.assertEquals(expectedOutput[0], (float) outLabelTensor.getValue(), 1e-3f); + } + } finally { + OnnxValue.close(pinnedInputs); + } + } + + @Test + public void TestTrainingSessionTrainStepOrtOutput() throws OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, null)) { + runTrainStep(trainingSession); + } + } + + @Test + public void TestSaveCheckpoint() throws IOException, OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + + Path tmpPath = Files.createTempDirectory("ort-java-training-test"); + try { + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, null)) { + + // Save checkpoint + trainingSession.saveCheckpoint(tmpPath, false); + } + + try (OrtTrainingSession trainingSession = + env.createTrainingSession(tmpPath.toString(), trainingPath, null, null)) { + // Load saved checkpoint into new session and run train step + runTrainStep(trainingSession); + } + } finally { + TestHelpers.deleteDirectoryTree(tmpPath); + } + } + + @Test + public void TestTrainingSessionOptimizerStep() throws OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + String optimizerPath = TestHelpers.getResourcePath("/adamw.onnx").toString(); + float[] expectedOutput_1 = + TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/loss_1.out")); + float[] expectedOutput_2 = + TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/loss_2.out")); + float[] input = TestHelpers.loadTensorFromFile(TestHelpers.getResourcePath("/input-0.in")); + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, optimizerPath)) { + int[] labels = {1, 1}; + + // Run train step with pinned inputs and pinned outputs + Map pinnedInputs = new HashMap<>(); + try { + // Create inputs + long[] inputShape = {2, 784}; + pinnedInputs.put( + "input-0", OnnxTensor.createTensor(env, OrtUtil.reshape(input, inputShape))); + + // long[] labelsShape = {2}; + pinnedInputs.put("labels", OnnxTensor.createTensor(env, labels)); + + try (OrtSession.Result outputs = trainingSession.trainStep(pinnedInputs)) { + Assertions.assertEquals(expectedOutput_1[0], (float) outputs.get(0).getValue(), 1e-3f); + } + + trainingSession.lazyResetGrad(); + + try (OrtSession.Result outputs = trainingSession.trainStep(pinnedInputs)) { + Assertions.assertEquals(expectedOutput_1[0], (float) outputs.get(0).getValue(), 1e-3f); + } + + trainingSession.optimizerStep(); + + try (OrtSession.Result outputs = trainingSession.trainStep(pinnedInputs)) { + Assertions.assertEquals(expectedOutput_2[0], (float) outputs.get(0).getValue(), 1e-3f); + } + } finally { + OnnxValue.close(pinnedInputs); + } + } + } + + @Test + public void TestTrainingSessionSetLearningRate() throws OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + String optimizerPath = TestHelpers.getResourcePath("/adamw.onnx").toString(); + + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, optimizerPath)) { + float learningRate = 0.245f; + trainingSession.setLearningRate(learningRate); + float actualLearningRate = trainingSession.getLearningRate(); + Assertions.assertEquals(learningRate, actualLearningRate); + } + } + + @Test + public void TestTrainingSessionLinearLRScheduler() throws OrtException { + String checkpointPath = TestHelpers.getResourcePath("/checkpoint.ckpt").toString(); + String trainingPath = TestHelpers.getResourcePath("/training_model.onnx").toString(); + String optimizerPath = TestHelpers.getResourcePath("/adamw.onnx").toString(); + + try (OrtTrainingSession trainingSession = + env.createTrainingSession(checkpointPath, trainingPath, null, optimizerPath)) { + float learningRate = 0.1f; + trainingSession.registerLinearLRScheduler(2, 4, learningRate); + runTrainStep(trainingSession); + trainingSession.optimizerStep(); + trainingSession.schedulerStep(); + Assertions.assertEquals(0.05f, trainingSession.getLearningRate()); + trainingSession.optimizerStep(); + trainingSession.schedulerStep(); + Assertions.assertEquals(0.1f, trainingSession.getLearningRate()); + trainingSession.optimizerStep(); + trainingSession.schedulerStep(); + Assertions.assertEquals(0.05f, trainingSession.getLearningRate()); + trainingSession.optimizerStep(); + trainingSession.schedulerStep(); + Assertions.assertEquals(0.0f, trainingSession.getLearningRate()); + } + } +} diff --git a/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h index 9416df69ee..ed4a20f7cd 100644 --- a/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h +++ b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h @@ -4,6 +4,7 @@ // This file contains the training c apis. #pragma once +#include #include "onnxruntime_c_api.h" ORT_RUNTIME_CLASS(TrainingSession); /// Type that enables performing training for the given user models. diff --git a/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-training-apis.yml b/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-training-apis.yml index e9eeb4d84a..6c0307eef0 100644 --- a/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-training-apis.yml +++ b/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-training-apis.yml @@ -22,6 +22,7 @@ jobs: --enable_training_apis \ --use_cuda --cuda_version=11.6 --cuda_home=/usr/local/cuda-11.6 --cudnn_home=/usr/local/cuda-11.6 \ --build_wheel \ + --build_java \ --skip_tests \ " \ -u