saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Java] Initial on device training support (#14027)

Browse source 
contributor: @Craigacp
This commit is contained in:
Adam Pocock 2023-03-08 13:01:08 -05:00 committed by GitHub
parent f14ab63c19
commit 47f00b5d49
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 2297 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
cmake/onnxruntime_java.cmake
View file

@ -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})

15
cmake/onnxruntime_java_unittests.cmake
View file

@ -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")

11
cmake/onnxruntime_unittests.cmake
View file

@ -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 $<TARGET_FILE:custom_op_library>
${JAVA_NATIVE_TEST_DIR}/$<TARGET_LINKER_FILE_NAME:custom_op_library>)
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()

18
java/build.gradle
View file

@ -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 {

33
java/src/main/java/ai/onnxruntime/OnnxRuntime.java
View file

@ -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<OrtProvider> 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.
*

93
java/src/main/java/ai/onnxruntime/OrtEnvironment.java
View file

@ -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.
*

47
java/src/main/java/ai/onnxruntime/OrtSession.java
View file

@ -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}.
*
* <p>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);
}

964
java/src/main/java/ai/onnxruntime/OrtTrainingSession.java Normal file
View file

@ -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.
*
* <p>Allows the inspection of the model's input and output nodes. Produced by an {@link
* OrtEnvironment}.
*
* <p>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<String> trainInputNames;
private final Set<String> trainOutputNames;
private final Set<String> evalInputNames;
private final Set<String> evalOutputNames;
private boolean closed = false;
/**
* Constructs an {@code OrtTrainingSession}.
*
* <p>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.
*
* <p>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).
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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<String> getTrainInputNames() {
return trainInputNames;
}
/**
* Returns an ordered set of the train model output names.
*
* @return The training outputs.
*/
public Set<String> getTrainOutputNames() {
return trainOutputNames;
}
/**
* Returns an ordered set of the eval model input names.
*
* @return The evaluation inputs.
*/
public Set<String> getEvalInputNames() {
return evalInputNames;
}
/**
* Returns an ordered set of the eval model output names.
*
* @return The evaluation outputs.
*/
public Set<String> 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.
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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}.
*
* <p>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.
*
* <p>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.
*
* <p>\param[in] session The training session which has working knowledge of the eval model.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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<String, ? extends OnnxTensorLike> 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<String, ? extends OnnxTensorLike> 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<String, ? extends OnnxTensorLike> inputs, Set<String> 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<String, ? extends OnnxTensorLike> inputs,
Set<String> 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<String, ? extends OnnxTensorLike> 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
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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<String, ? extends OnnxTensorLike> 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<String, ? extends OnnxTensorLike> 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<String, ? extends OnnxTensorLike> inputs, Set<String> 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<String, ? extends OnnxTensorLike> inputs,
Set<String> 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<String, ? extends OnnxTensorLike> 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
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>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`.
*
* <p>\param[in] sess The training session on which the learning rate needs to be set. \param[in]
* learning_rate Desired learning rate to set.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>This function allows users to get the learning rate for the training session. The current
* learning rate is maintained by the training session
*
* <p>\param[in] sess The training session on which the learning rate needs to be set.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>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.
*
* <p>\param[in] sess The training session which has working knowledge of the optimizer model.
* \param[in] run_options Run options for this eval step.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>\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.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>\param[in] sess The training session that has the registered learning rate scheduler.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>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.
*
* <p>\param[in] sess The training session. \param[in] inference_model_path Path where the
* inference model should be serialized to.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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.
*
* <p>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.
*
* <p>\param[in] checkpoint_path Path to the checkpoint directory \param[out] checkpoint_state
* Checkpoint states that contains the states of the training session.
*
* <p>\snippet{doc} snippets.dox OrtStatus Return Value
*
* <p>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);
}
}

7
java/src/main/java/ai/onnxruntime/providers/package-info.java Normal file
View file

@ -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;

14
java/src/main/native/ai_onnxruntime_OnnxRuntime.c
View file

@ -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;
}
/*

16
java/src/main/native/ai_onnxruntime_OrtSession_RunOptions.c
View file

@ -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

762
java/src/main/native/ai_onnxruntime_OrtTrainingSession.c Normal file
View file

@ -0,0 +1,762 @@
/*
* Copyright (c) 2022 Oracle and/or its affiliates. All rights reserved.
* Licensed under the MIT License.
*/
#include <jni.h>
#include <string.h>
#include <stdlib.h>
#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

58
java/src/main/native/ai_onnxruntime_OrtTrainingSession_OrtCheckpointState.c Normal file
View file

@ -0,0 +1,58 @@
/*
* Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
* Licensed under the MIT License.
*/
#include <jni.h>
#include <string.h>
#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);
}

4
java/src/test/java/ai/onnxruntime/InferenceTest.java
View file

@ -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.

9
java/src/test/java/ai/onnxruntime/SparseTensorTest.java
View file

@ -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)) {

8
java/src/test/java/ai/onnxruntime/TestHelpers.java
View file

@ -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<Boolean> 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) {

256
java/src/test/java/ai/onnxruntime/TrainingTest.java Normal file
View file

@ -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<String> inputNames = trainingSession.getTrainInputNames();
Assertions.assertFalse(inputNames.isEmpty());
Set<String> 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<String, OnnxTensor> pinnedInputs = new HashMap<>();
Map<String, OnnxTensor> 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<String, OnnxTensor> 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<String, OnnxTensor> 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());
}
}
}

1
orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h
View file

@ -4,6 +4,7 @@
// This file contains the training c apis.
#pragma once
#include <stdbool.h>
#include "onnxruntime_c_api.h"
ORT_RUNTIME_CLASS(TrainingSession); /// Type that enables performing training for the given user models.

1
tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-training-apis.yml
View file

@ -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