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

[java] Enable output pinning in OrtSession and OrtTrainingSession (#16835)

Browse source
This commit is contained in:
Adam Pocock 2023-09-26 04:49:13 -04:00 committed by GitHub
parent ccb73fd827
commit aed43f429a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 668 additions and 119 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -239,7 +239,7 @@ public class OrtSession implements AutoCloseable {
*/
public Result run(Map<String, ? extends OnnxTensorLike> inputs, Set<String> requestedOutputs)
throws OrtException {
return run(inputs, requestedOutputs, null);
return run(inputs, requestedOutputs, Collections.emptyMap(), null);
}
/**
@ -259,17 +259,90 @@ public class OrtSession implements AutoCloseable {
Set<String> requestedOutputs,
RunOptions runOptions)
throws OrtException {
return run(inputs, requestedOutputs, Collections.emptyMap(), runOptions);
}
/**
* Scores an input feed dict, returning the map of pinned outputs.
*
* <p>The outputs are sorted based on the supplied map traversal order.
*
* <p>Note: pinned outputs are not owned by the {@link Result} object, and are <b>not</b> closed
* when the result object is closed.
*
* @param inputs The inputs to score.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @return The inferred outputs.
* @throws OrtException If there was an error in native code, the input or output names are
* invalid, or if there are zero or too many inputs or outputs.
*/
public Result run(
Map<String, ? extends OnnxTensorLike> inputs, Map<String, ? extends OnnxValue> pinnedOutputs)
throws OrtException {
return run(inputs, Collections.emptySet(), pinnedOutputs, null);
}
/**
* Scores an input feed dict, returning the map of requested and pinned outputs.
*
* <p>The outputs are sorted based on the supplied set traversal order with pinned outputs first,
* then requested outputs. An {@link IllegalArgumentException} is thrown if the same output name
* appears in both the requested outputs and the pinned outputs.
*
* <p>Note: pinned outputs are not owned by the {@link Result} object, and are <b>not</b> closed
* when the result object is closed.
*
* @param inputs The inputs to score.
* @param requestedOutputs The requested outputs which ORT will allocate.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @return The inferred outputs.
* @throws OrtException If there was an error in native code, the input or output names are
* invalid, or if there are zero or too many inputs or outputs.
*/
public Result run(
Map<String, ? extends OnnxTensorLike> inputs,
Set<String> requestedOutputs,
Map<String, ? extends OnnxValue> pinnedOutputs)
throws OrtException {
return run(inputs, requestedOutputs, pinnedOutputs, null);
}
/**
* Scores an input feed dict, returning the map of requested and pinned outputs.
*
* <p>The outputs are sorted based on the supplied set traversal order with pinned outputs first,
* then requested outputs. An {@link IllegalArgumentException} is thrown if the same output name
* appears in both the requested outputs and the pinned outputs.
*
* <p>Note: pinned outputs are not owned by the {@link Result} object, and are <b>not</b> closed
* when the result object is closed.
*
* @param inputs The inputs to score.
* @param requestedOutputs The requested outputs which ORT will allocate.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @param runOptions The RunOptions to control this run.
* @return The inferred outputs.
* @throws OrtException If there was an error in native code, the input or output names are
* invalid, or if there are zero or too many inputs or outputs.
*/
public Result run(
Map<String, ? extends OnnxTensorLike> inputs,
Set<String> requestedOutputs,
Map<String, ? extends OnnxValue> pinnedOutputs,
RunOptions runOptions)
throws OrtException {
if (!closed) {
if ((inputs.isEmpty() && (numInputs != 0)) || (inputs.size() > numInputs)) {
throw new OrtException(
"Unexpected number of inputs, expected [1," + numInputs + ") found " + inputs.size());
}
if (requestedOutputs.isEmpty() || (requestedOutputs.size() > numOutputs)) {
int totalOutputs = requestedOutputs.size() + pinnedOutputs.size();
if ((totalOutputs == 0) || (totalOutputs > numOutputs)) {
throw new OrtException(
"Unexpected number of requestedOutputs, expected [1,"
"Unexpected number of requestedOutputs & pinnedOutputs, expected [1,"
+ numOutputs
+ ") found "
+ requestedOutputs.size());
+ totalOutputs);
}
String[] inputNamesArray = new String[inputs.size()];
long[] inputHandles = new long[inputs.size()];
@ -284,20 +357,41 @@ public class OrtSession implements AutoCloseable {
"Unknown input name " + t.getKey() + ", expected one of " + inputNames.toString());
}
}
String[] outputNamesArray = new String[requestedOutputs.size()];
String[] outputNamesArray = new String[requestedOutputs.size() + pinnedOutputs.size()];
OnnxValue[] outputValues = new OnnxValue[outputNamesArray.length];
long[] outputHandles = new long[outputNamesArray.length];
i = 0;
for (Map.Entry<String, ? extends OnnxValue> e : pinnedOutputs.entrySet()) {
if (outputNames.contains(e.getKey())) {
outputNamesArray[i] = e.getKey();
outputValues[i] = e.getValue();
outputHandles[i] = getHandle(e.getValue());
i++;
} else {
throw new OrtException(
"Unknown output name " + e.getKey() + ", expected one of " + outputNames.toString());
}
}
for (String s : requestedOutputs) {
if (outputNames.contains(s)) {
outputNamesArray[i] = s;
i++;
if (!pinnedOutputs.containsKey(s)) {
outputNamesArray[i] = s;
// outputValues and outputHandles can be null/0 for these outputs as ORT will allocate
// them.
i++;
} else {
throw new OrtException(
"Output '"
+ s
+ "' was found in both the requested outputs and the pinned outputs");
}
} else {
throw new OrtException(
"Unknown output name " + s + ", expected one of " + outputNames.toString());
}
}
long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle();
OnnxValue[] outputValues =
boolean[] ownedByResult =
run(
OnnxRuntime.ortApiHandle,
nativeHandle,
@ -307,13 +401,40 @@ public class OrtSession implements AutoCloseable {
inputNamesArray.length,
outputNamesArray,
outputNamesArray.length,
outputValues,
outputHandles,
runOptionsHandle);
return new Result(outputNamesArray, outputValues);
return new Result(outputNamesArray, outputValues, ownedByResult);
} else {
throw new IllegalStateException("Trying to score a closed OrtSession.");
}
}
/**
* Pulls out the native handle by casting it to the appropriate type.
*
* @param v The OnnxValue.
* @return The native handle.
*/
static long getHandle(OnnxValue v) {
/*
* Note this method exists as interface methods are all public, but we do not want users to be
* able to access the native pointer via a public API so can't add a method to OnnxValue which
* exposes it.
*/
if (v instanceof OnnxTensorLike) {
return ((OnnxTensorLike) v).nativeHandle;
} else if (v instanceof OnnxSequence) {
return ((OnnxSequence) v).nativeHandle;
} else if (v instanceof OnnxMap) {
return ((OnnxMap) v).nativeHandle;
} else {
throw new IllegalArgumentException(
"Unexpected OnnxValue subclass, should be {OnnxTensorLike, OnnxSequence, OnnxMap}, found "
+ v.getClass());
}
}
/**
* Gets the metadata for the currently loaded model.
*
@ -409,8 +530,9 @@ public class OrtSession implements AutoCloseable {
throws OrtException;
/**
* The native run call. runOptionsHandle can be zero (i.e. the null pointer), but all other
* handles must be valid pointers.
* The native run call. runOptionsHandle can be zero (i.e. the null pointer), outputValues can
* contain null entries, and outputHandles can contain zero values (i.e. the null pointer), but
* all other handles must be valid pointers.
*
* @param apiHandle The pointer to the api.
* @param nativeHandle The pointer to the session.
@ -419,12 +541,14 @@ public class OrtSession implements AutoCloseable {
* @param inputs The input tensors.
* @param numInputs The number of inputs.
* @param outputNamesArray The requested output names.
* @param outputValues The OnnxValue output array.
* @param outputHandles The OrtValue output pointer array.
* @param numOutputs The number of requested outputs.
* @param runOptionsHandle The (possibly null) pointer to the run options.
* @return The OnnxValues produced by this run.
* @return A boolean array representing if the OnnxValues were allocated by this run call.
* @throws OrtException If the native call failed in some way.
*/
private native OnnxValue[] run(
private native boolean[] run(
long apiHandle,
long nativeHandle,
long allocatorHandle,
@ -433,6 +557,8 @@ public class OrtSession implements AutoCloseable {
long numInputs,
String[] outputNamesArray,
long numOutputs,
OnnxValue[] outputValues,
long[] outputHandles,
long runOptionsHandle)
throws OrtException;
@ -1417,9 +1543,13 @@ public class OrtSession implements AutoCloseable {
/**
* An {@link AutoCloseable} wrapper around a {@link Map} containing {@link OnnxValue}s.
*
* <p>When this is closed it closes all the {@link OnnxValue}s inside it. If you maintain a
* reference to a value after this object has been closed it will throw an {@link
* <p>When this is closed it closes all the {@link OnnxValue}s owned by the result object. If you
* maintain a reference to a value after this object has been closed it will throw an {@link
* IllegalStateException} upon access.
*
* <p>{@link OnnxValue}s which are supplied as pinned outputs to a {@code run} call are not closed
* by the {@link Result#close()} method. Ownership of each output can be checked with {@link
* Result#isResultOwner(int)}.
*/
public static class Result implements AutoCloseable, Iterable<Map.Entry<String, OnnxValue>> {
@ -1429,6 +1559,8 @@ public class OrtSession implements AutoCloseable {
private final List<OnnxValue> list;
private final boolean[] ownedByResult;
private boolean closed;
/**
@ -1437,21 +1569,23 @@ public class OrtSession implements AutoCloseable {
* @param names The output names.
* @param values The output values.
*/
Result(String[] names, OnnxValue[] values) {
if (names.length != values.length) {
Result(String[] names, OnnxValue[] values, boolean[] ownedByResult) {
if ((names.length != values.length) || (names.length != ownedByResult.length)) {
throw new IllegalArgumentException(
"Expected same number of names and values, found names.length = "
"Expected same number of names, values and ownedByResult, found names.length = "
+ names.length
+ ", values.length = "
+ values.length);
+ values.length
+ ", ownedByResult.length = "
+ ownedByResult.length);
}
map = new LinkedHashMap<>(OrtUtil.capacityFromSize(names.length));
list = new ArrayList<>(names.length);
list = new ArrayList<>(Arrays.asList(values));
this.ownedByResult = ownedByResult;
for (int i = 0; i < names.length; i++) {
map.put(names[i], values[i]);
list.add(values[i]);
}
this.closed = false;
}
@ -1460,8 +1594,11 @@ public class OrtSession implements AutoCloseable {
public void close() {
if (!closed) {
closed = true;
for (OnnxValue t : map.values()) {
t.close();
for (int i = 0; i < list.size(); i++) {
if (ownedByResult[i]) {
OnnxValue value = list.get(i);
value.close();
}
}
} else {
logger.warning("Closing an already closed Result");
@ -1494,6 +1631,23 @@ public class OrtSession implements AutoCloseable {
}
}
/**
* Gets the value from the container at the specified index.
*
* <p>Throws {@link IllegalStateException} if the container has been closed, and {@link
* ArrayIndexOutOfBoundsException} if the index is invalid.
*
* @param index The index to lookup.
* @return Is that value owned by this result object?
*/
public boolean isResultOwner(int index) {
if (!closed) {
return ownedByResult[index];
} else {
throw new IllegalStateException("Result is closed");
}
}
/**
* Returns the number of outputs in this Result.
*

172
java/src/main/java/ai/onnxruntime/OrtTrainingSession.java
View file

@ -418,7 +418,7 @@ public final class OrtTrainingSession implements AutoCloseable {
*/
public OrtSession.Result trainStep(Map<String, ? extends OnnxTensorLike> inputs)
throws OrtException {
return trainStep(inputs, trainOutputNames, null);
return trainStep(inputs, trainOutputNames, Collections.emptyMap(), null);
}
/**
@ -432,7 +432,7 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result trainStep(
Map<String, ? extends OnnxTensorLike> inputs, OrtSession.RunOptions runOptions)
throws OrtException {
return trainStep(inputs, trainOutputNames, runOptions);
return trainStep(inputs, trainOutputNames, Collections.emptyMap(), runOptions);
}
/**
@ -446,14 +446,41 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result trainStep(
Map<String, ? extends OnnxTensorLike> inputs, Set<String> requestedOutputs)
throws OrtException {
return trainStep(inputs, requestedOutputs, null);
return trainStep(inputs, requestedOutputs, Collections.emptyMap(), null);
}
/**
* Performs a single step of training, accumulating the gradients.
*
* <p>The outputs are sorted based on the supplied map traversal order.
*
* <p>Note: pinned outputs are not owned by the {@link OrtSession.Result} object, and are
* <b>not</b> closed when the result object is closed.
*
* @param inputs The inputs (must include both the features and the target).
* @param requestedOutputs The requested outputs.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @return Requested outputs produced by the training step.
* @throws OrtException If the native call failed.
*/
public OrtSession.Result trainStep(
Map<String, ? extends OnnxTensorLike> inputs, Map<String, ? extends OnnxValue> pinnedOutputs)
throws OrtException {
return trainStep(inputs, Collections.emptySet(), pinnedOutputs, null);
}
/**
* Performs a single step of training, accumulating the gradients.
*
* <p>The outputs are sorted based on the supplied set traversal order with pinned outputs first,
* then requested outputs. An {@link IllegalArgumentException} is thrown if the same output name
* appears in both the requested outputs and the pinned outputs.
*
* <p>Note: pinned outputs are not owned by the {@link OrtSession.Result} object, and are
* <b>not</b> closed when the result object is closed.
*
* @param inputs The inputs (must include both the features and the target).
* @param requestedOutputs The requested outputs which ORT will allocate.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @param runOptions Run options for controlling this specific call.
* @return Requested outputs produced by the training step.
* @throws OrtException If the native call failed.
@ -461,6 +488,7 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result trainStep(
Map<String, ? extends OnnxTensorLike> inputs,
Set<String> requestedOutputs,
Map<String, ? extends OnnxValue> pinnedOutputs,
OrtSession.RunOptions runOptions)
throws OrtException {
checkClosed();
@ -472,12 +500,14 @@ public final class OrtTrainingSession implements AutoCloseable {
+ ") found "
+ inputs.size());
}
if (requestedOutputs.isEmpty() || (requestedOutputs.size() > trainOutputNames.size())) {
int numTrainOutputs = trainOutputNames.size();
int totalOutputs = requestedOutputs.size() + pinnedOutputs.size();
if ((totalOutputs == 0) || (totalOutputs > numTrainOutputs)) {
throw new OrtException(
"Unexpected number of requestedOutputs, expected [1,"
+ trainOutputNames.size()
"Unexpected number of requestedOutputs & pinnedOutputs, expected [1,"
+ numTrainOutputs
+ ") found "
+ requestedOutputs.size());
+ totalOutputs);
}
String[] inputNamesArray = new String[inputs.size()];
long[] inputHandles = new long[inputs.size()];
@ -492,12 +522,35 @@ public final class OrtTrainingSession implements AutoCloseable {
"Unknown input name " + t.getKey() + ", expected one of " + trainInputNames);
}
}
String[] outputNamesArray = new String[requestedOutputs.size()];
String[] outputNamesArray = new String[requestedOutputs.size() + pinnedOutputs.size()];
OnnxValue[] outputValues = new OnnxValue[outputNamesArray.length];
long[] outputHandles = new long[outputNamesArray.length];
i = 0;
for (Map.Entry<String, ? extends OnnxValue> e : pinnedOutputs.entrySet()) {
if (trainOutputNames.contains(e.getKey())) {
outputNamesArray[i] = e.getKey();
outputValues[i] = e.getValue();
outputHandles[i] = OrtSession.getHandle(e.getValue());
i++;
} else {
throw new OrtException(
"Unknown output name "
+ e.getKey()
+ ", expected one of "
+ trainOutputNames.toString());
}
}
for (String s : requestedOutputs) {
if (trainOutputNames.contains(s)) {
outputNamesArray[i] = s;
i++;
if (!pinnedOutputs.containsKey(s)) {
outputNamesArray[i] = s;
// outputValues and outputHandles can be null/0 for these outputs as ORT will allocate
// them.
i++;
} else {
throw new OrtException(
"Output '" + s + "' was found in both the requested outputs and the pinned outputs");
}
} else {
throw new OrtException(
"Unknown output name " + s + ", expected one of " + trainOutputNames.toString());
@ -505,7 +558,7 @@ public final class OrtTrainingSession implements AutoCloseable {
}
long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle();
OnnxValue[] outputValues =
boolean[] ownedByResult =
trainStep(
OnnxRuntime.ortApiHandle,
OnnxRuntime.ortTrainingApiHandle,
@ -516,8 +569,10 @@ public final class OrtTrainingSession implements AutoCloseable {
inputNamesArray.length,
outputNamesArray,
outputNamesArray.length,
outputValues,
outputHandles,
runOptionsHandle);
return new OrtSession.Result(outputNamesArray, outputValues);
return new OrtSession.Result(outputNamesArray, outputValues, ownedByResult);
}
/*
@ -540,7 +595,7 @@ public final class OrtTrainingSession implements AutoCloseable {
* 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(
private native boolean[] trainStep(
long apiHandle,
long trainingApiHandle,
long nativeHandle,
@ -550,6 +605,8 @@ public final class OrtTrainingSession implements AutoCloseable {
long numInputs,
String[] outputNamesArray,
long numOutputs,
OnnxValue[] outputValues,
long[] outputHandles,
long runOptionsHandle);
/**
@ -561,7 +618,7 @@ public final class OrtTrainingSession implements AutoCloseable {
*/
public OrtSession.Result evalStep(Map<String, ? extends OnnxTensorLike> inputs)
throws OrtException {
return evalStep(inputs, evalOutputNames, null);
return evalStep(inputs, evalOutputNames, Collections.emptyMap(), null);
}
/**
@ -575,7 +632,7 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result evalStep(
Map<String, ? extends OnnxTensorLike> inputs, OrtSession.RunOptions runOptions)
throws OrtException {
return evalStep(inputs, evalOutputNames, runOptions);
return evalStep(inputs, evalOutputNames, Collections.emptyMap(), runOptions);
}
/**
@ -589,14 +646,41 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result evalStep(
Map<String, ? extends OnnxTensorLike> inputs, Set<String> requestedOutputs)
throws OrtException {
return evalStep(inputs, requestedOutputs, null);
return evalStep(inputs, requestedOutputs, Collections.emptyMap(), null);
}
/**
* Performs a single evaluation step using the supplied inputs.
*
* @param inputs The model inputs.
* @param requestedOutputs The requested output names.
* <p>The outputs are sorted based on the supplied map traversal order.
*
* <p>Note: pinned outputs are not owned by the {@link OrtSession.Result} object, and are
* <b>not</b> closed when the result object is closed.
*
* @param inputs The inputs to score.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @return The requested outputs.
* @throws OrtException If the native call failed.
*/
public OrtSession.Result evalStep(
Map<String, ? extends OnnxTensorLike> inputs, Map<String, ? extends OnnxValue> pinnedOutputs)
throws OrtException {
return evalStep(inputs, Collections.emptySet(), pinnedOutputs, null);
}
/**
* Performs a single evaluation step using the supplied inputs.
*
* <p>The outputs are sorted based on the supplied set traversal order with pinned outputs first,
* then requested outputs. An {@link IllegalArgumentException} is thrown if the same output name
* appears in both the requested outputs and the pinned outputs.
*
* <p>Note: pinned outputs are not owned by the {@link OrtSession.Result} object, and are
* <b>not</b> closed when the result object is closed.
*
* @param inputs The inputs to score.
* @param requestedOutputs The requested outputs which ORT will allocate.
* @param pinnedOutputs The requested outputs which the user has allocated.
* @param runOptions Run options for controlling this specific call.
* @return The requested outputs.
* @throws OrtException If the native call failed.
@ -604,6 +688,7 @@ public final class OrtTrainingSession implements AutoCloseable {
public OrtSession.Result evalStep(
Map<String, ? extends OnnxTensorLike> inputs,
Set<String> requestedOutputs,
Map<String, ? extends OnnxValue> pinnedOutputs,
OrtSession.RunOptions runOptions)
throws OrtException {
checkClosed();
@ -615,12 +700,14 @@ public final class OrtTrainingSession implements AutoCloseable {
+ ") found "
+ inputs.size());
}
if (requestedOutputs.isEmpty() || (requestedOutputs.size() > evalOutputNames.size())) {
int numEvalOutputs = evalOutputNames.size();
int totalOutputs = requestedOutputs.size() + pinnedOutputs.size();
if ((totalOutputs == 0) || (totalOutputs > numEvalOutputs)) {
throw new OrtException(
"Unexpected number of requestedOutputs, expected [1,"
+ evalOutputNames.size()
"Unexpected number of requestedOutputs & pinnedOutputs, expected [1,"
+ numEvalOutputs
+ ") found "
+ requestedOutputs.size());
+ totalOutputs);
}
String[] inputNamesArray = new String[inputs.size()];
long[] inputHandles = new long[inputs.size()];
@ -635,12 +722,35 @@ public final class OrtTrainingSession implements AutoCloseable {
"Unknown input name " + t.getKey() + ", expected one of " + evalInputNames.toString());
}
}
String[] outputNamesArray = new String[requestedOutputs.size()];
String[] outputNamesArray = new String[requestedOutputs.size() + pinnedOutputs.size()];
OnnxValue[] outputValues = new OnnxValue[outputNamesArray.length];
long[] outputHandles = new long[outputNamesArray.length];
i = 0;
for (Map.Entry<String, ? extends OnnxValue> e : pinnedOutputs.entrySet()) {
if (evalOutputNames.contains(e.getKey())) {
outputNamesArray[i] = e.getKey();
outputValues[i] = e.getValue();
outputHandles[i] = OrtSession.getHandle(e.getValue());
i++;
} else {
throw new OrtException(
"Unknown output name "
+ e.getKey()
+ ", expected one of "
+ evalOutputNames.toString());
}
}
for (String s : requestedOutputs) {
if (evalOutputNames.contains(s)) {
outputNamesArray[i] = s;
i++;
if (!pinnedOutputs.containsKey(s)) {
outputNamesArray[i] = s;
// outputValues and outputHandles can be null/0 for these outputs as ORT will allocate
// them.
i++;
} else {
throw new OrtException(
"Output '" + s + "' was found in both the requested outputs and the pinned outputs");
}
} else {
throw new OrtException(
"Unknown output name " + s + ", expected one of " + evalOutputNames.toString());
@ -648,7 +758,7 @@ public final class OrtTrainingSession implements AutoCloseable {
}
long runOptionsHandle = runOptions == null ? 0 : runOptions.getNativeHandle();
OnnxValue[] outputValues =
boolean[] ownedByResult =
evalStep(
OnnxRuntime.ortApiHandle,
OnnxRuntime.ortTrainingApiHandle,
@ -659,8 +769,10 @@ public final class OrtTrainingSession implements AutoCloseable {
inputNamesArray.length,
outputNamesArray,
outputNamesArray.length,
outputValues,
outputHandles,
runOptionsHandle);
return new OrtSession.Result(outputNamesArray, outputValues);
return new OrtSession.Result(outputNamesArray, outputValues, ownedByResult);
}
/*
@ -682,7 +794,7 @@ public final class OrtTrainingSession implements AutoCloseable {
* 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(
private native boolean[] evalStep(
long apiHandle,
long trainingApiHandle,
long nativeHandle,
@ -692,6 +804,8 @@ public final class OrtTrainingSession implements AutoCloseable {
long numInputs,
String[] outputNamesArray,
long numOutputs,
OnnxValue[] outputValues,
long[] outputHandles,
long runOptionsHandle)
throws OrtException;

38
java/src/main/native/ai_onnxruntime_OrtSession.c
View file

@ -316,14 +316,19 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtSession_getOutputInfo(JNIE
/*
* Class: ai_onnxruntime_OrtSession
* Method: run
* Signature: (JJJ[Ljava/lang/String;[JJ[Ljava/lang/String;JJ)[Lai/onnxruntime/OnnxValue;
* private native OnnxValue[] run(long apiHandle, long nativeHandle, long allocatorHandle, String[] inputNamesArray, long[] inputs, long numInputs, String[] outputNamesArray, long numOutputs)
* Signature: (JJJ[Ljava/lang/String;[JJ[Ljava/lang/String;J[Lai/onnxruntime/OnnxValue;[JJ)[Z
* private native boolean[] run(long apiHandle, long nativeHandle, long allocatorHandle,
* String[] inputNamesArray, long[] inputs, long numInputs,
* String[] outputNamesArray, long numOutputs,
* OnnxValue[] outputValues, long[] outputHandles,
* long runOptionsHandle) throws OrtException;
*/
JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtSession_run(JNIEnv* jniEnv, jobject jobj, jlong apiHandle,
JNIEXPORT jbooleanArray JNICALL Java_ai_onnxruntime_OrtSession_run(JNIEnv* jniEnv, jobject jobj, jlong apiHandle,
jlong sessionHandle, jlong allocatorHandle,
jobjectArray inputNamesArr, jlongArray tensorArr,
jlong numInputs, jobjectArray outputNamesArr,
jlong numOutputs, jlong runOptionsHandle) {
jlong numOutputs, jobjectArray outputValuesArr,
jlongArray outputHandlesArr, 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;
@ -331,7 +336,7 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtSession_run(JNIEnv* jniEnv
OrtSession* session = (OrtSession*)sessionHandle;
OrtRunOptions* runOptions = (OrtRunOptions*)runOptionsHandle;
jobjectArray outputArray = NULL;
jbooleanArray outputArray = NULL;
// Create the buffers for the Java input & output strings, and the input pointers
const char** inputNames = allocarray(numInputs, sizeof(char*));
@ -376,13 +381,19 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtSession_run(JNIEnv* jniEnv
// Release the java array copy of pointers to the tensors.
(*jniEnv)->ReleaseLongArrayElements(jniEnv, tensorArr, inputValueLongs, JNI_ABORT);
// Extract a C array of longs which are pointers to the output tensors.
jlong* outputHandleLongs = (*jniEnv)->GetLongArrayElements(jniEnv, outputHandlesArr, NULL);
// 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;
outputValues[i] = (OrtValue*)outputHandleLongs[i];
}
// Release the java array copy of pointers to the outputs.
(*jniEnv)->ReleaseLongArrayElements(jniEnv, outputHandlesArr, outputHandleLongs, JNI_ABORT);
// Actually score the inputs.
// ORT_API_STATUS(OrtRun, _Inout_ OrtSession* sess, _In_ OrtRunOptions* run_options,
// _In_ const char* const* input_names, _In_ const OrtValue* const* input, size_t input_len,
@ -394,21 +405,26 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtSession_run(JNIEnv* jniEnv
goto cleanup_output_values;
}
// Construct the output array of ONNXValues
jclass onnxValueClass = (*jniEnv)->FindClass(jniEnv, ORTJNI_OnnxValueClassName);
outputArray = (*jniEnv)->NewObjectArray(jniEnv, safecast_int64_to_jsize(numOutputs), onnxValueClass, NULL);
// Create the output boolean array denoting if ORT owns the memory for each output.
// Java boolean arrays are initialized to false.
outputArray = (*jniEnv)->NewBooleanArray(jniEnv, safecast_int64_to_jsize(numOutputs));
jboolean* boolArr = (*jniEnv)->GetBooleanArrayElements(jniEnv, outputArray, NULL);
// Convert the output tensors into ONNXValues
for (int i = 0; i < numOutputs; i++) {
if (outputValues[i] != NULL) {
if (outputValues[i] != NULL && (*jniEnv)->GetObjectArrayElement(jniEnv, outputValuesArr, 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);
boolArr[i] = 1;
(*jniEnv)->SetObjectArrayElement(jniEnv, outputValuesArr, i, onnxValue);
}
}
// Write the output array back to Java.
(*jniEnv)->ReleaseBooleanArrayElements(jniEnv, outputArray, boolArr, 0);
// 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:

109
java/src/main/native/ai_onnxruntime_OrtTrainingSession.c
View file

@ -1,5 +1,5 @@
/*
* Copyright (c) 2022 Oracle and/or its affiliates. All rights reserved.
* Copyright (c) 2022, 2023, Oracle and/or its affiliates. All rights reserved.
* Licensed under the MIT License.
*/
#include <jni.h>
@ -330,12 +330,12 @@ JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtTrainingSession_lazyResetGrad
/*
* Class: ai_onnxruntime_OrtTrainingSession
* Method: trainStep
* Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;JJ)[Lai/onnxruntime/OnnxValue;
* Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;J[Lai/onnxruntime/OnnxValue;[JJ)[Z
*/
JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
JNIEXPORT jbooleanArray 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) {
jobjectArray outputNamesArr, jlong numOutputs, jobjectArray outputValuesArr, jlongArray outputHandlesArr, 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;
@ -343,31 +343,31 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle;
OrtRunOptions* runOptions = (OrtRunOptions*)runOptionsHandle;
jobjectArray outputArray = NULL;
jbooleanArray outputArray = NULL;
// Create the buffers for the Java input & output strings, and the input pointers
const char** inputNames = malloc(sizeof(char*) * numInputs);
const char** inputNames = allocarray(numInputs, sizeof(char*));
if (inputNames == NULL) {
// Nothing to cleanup, return and throw exception
return outputArray;
}
const char** outputNames = malloc(sizeof(char*) * numOutputs);
const char** outputNames = allocarray(numOutputs, sizeof(char*));
if (outputNames == NULL) {
goto cleanup_input_names;
}
jobject* javaInputStrings = malloc(sizeof(jobject) * numInputs);
jobject* javaInputStrings = allocarray(numInputs, sizeof(jobject));
if (javaInputStrings == NULL) {
goto cleanup_output_names;
}
jobject* javaOutputStrings = malloc(sizeof(jobject) * numOutputs);
jobject* javaOutputStrings = allocarray(numOutputs, sizeof(jobject));
if (javaOutputStrings == NULL) {
goto cleanup_java_input_strings;
}
const OrtValue** inputValuePtrs = malloc(sizeof(OrtValue*) * numInputs);
const OrtValue** inputValuePtrs = allocarray(numInputs, sizeof(OrtValue*));
if (inputValuePtrs == NULL) {
goto cleanup_java_output_strings;
}
OrtValue** outputValues = malloc(sizeof(OrtValue*) * numOutputs);
OrtValue** outputValues = allocarray(numOutputs, sizeof(OrtValue*));
if (outputValues == NULL) {
goto cleanup_input_values;
}
@ -388,13 +388,19 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
// Release the java array copy of pointers to the tensors.
(*jniEnv)->ReleaseLongArrayElements(jniEnv, inputHandles, inputValueLongs, JNI_ABORT);
// Extract a C array of longs which are pointers to the output tensors.
jlong* outputHandleLongs = (*jniEnv)->GetLongArrayElements(jniEnv, outputHandlesArr, NULL);
// 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;
outputValues[i] = (OrtValue*)outputHandleLongs[i];
}
// Release the java array copy of pointers to the outputs.
(*jniEnv)->ReleaseLongArrayElements(jniEnv, outputHandlesArr, outputHandleLongs, JNI_ABORT);
// 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,
@ -406,24 +412,29 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
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);
// Create the output boolean array denoting if ORT owns the memory for each output.
// Java boolean arrays are initialized to false.
outputArray = (*jniEnv)->NewBooleanArray(jniEnv, safecast_int64_to_jsize(numOutputs));
jboolean* boolArr = (*jniEnv)->GetBooleanArrayElements(jniEnv, outputArray, NULL);
// Convert the output tensors into ONNXValues
for (int i = 0; i < numOutputs; i++) {
if (outputValues[i] != NULL) {
if (outputValues[i] != NULL && (*jniEnv)->GetObjectArrayElement(jniEnv, outputValuesArr, 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);
boolArr[i] = 1;
(*jniEnv)->SetObjectArrayElement(jniEnv, outputValuesArr, i, onnxValue);
}
}
// Write the output array back to Java.
(*jniEnv)->ReleaseBooleanArrayElements(jniEnv, outputArray, boolArr, 0);
// 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:
cleanup_output_values:
free(outputValues);
// Release the Java output strings
@ -437,15 +448,15 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
}
// Release the buffers
cleanup_input_values:
cleanup_input_values:
free((void*)inputValuePtrs);
cleanup_java_output_strings:
cleanup_java_output_strings:
free(javaOutputStrings);
cleanup_java_input_strings:
cleanup_java_input_strings:
free(javaInputStrings);
cleanup_output_names:
cleanup_output_names:
free((void*)outputNames);
cleanup_input_names:
cleanup_input_names:
free((void*)inputNames);
return outputArray;
@ -454,12 +465,12 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_trainStep
/*
* Class: ai_onnxruntime_OrtTrainingSession
* Method: evalStep
* Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;JJ)[Lai/onnxruntime/OnnxValue;
* Signature: (JJJJ[Ljava/lang/String;[JJ[Ljava/lang/String;J[Lai/onnxruntime/OnnxValue;[JJ)[Z
*/
JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep
JNIEXPORT jbooleanArray 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) {
jobjectArray outputNamesArr, jlong numOutputs, jobjectArray outputValuesArr, jlongArray outputHandlesArr, 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;
@ -467,31 +478,31 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep
OrtTrainingSession* trainSession = (OrtTrainingSession*)nativeHandle;
OrtRunOptions* runOptions = (OrtRunOptions*)runOptionsHandle;
jobjectArray outputArray = NULL;
jbooleanArray outputArray = NULL;
// Create the buffers for the Java input & output strings, and the input pointers
const char** inputNames = malloc(sizeof(char*) * numInputs);
const char** inputNames = allocarray(numInputs, sizeof(char*));
if (inputNames == NULL) {
// Nothing to cleanup, return and throw exception
return outputArray;
}
const char** outputNames = malloc(sizeof(char*) * numOutputs);
const char** outputNames = allocarray(numOutputs, sizeof(char*));
if (outputNames == NULL) {
goto cleanup_input_names;
}
jobject* javaInputStrings = malloc(sizeof(jobject) * numInputs);
jobject* javaInputStrings = allocarray(numInputs, sizeof(jobject));
if (javaInputStrings == NULL) {
goto cleanup_output_names;
}
jobject* javaOutputStrings = malloc(sizeof(jobject) * numOutputs);
jobject* javaOutputStrings = allocarray(numOutputs, sizeof(jobject));
if (javaOutputStrings == NULL) {
goto cleanup_java_input_strings;
}
const OrtValue** inputValuePtrs = malloc(sizeof(OrtValue*) * numInputs);
const OrtValue** inputValuePtrs = allocarray(numInputs, sizeof(OrtValue*));
if (inputValuePtrs == NULL) {
goto cleanup_java_output_strings;
}
OrtValue** outputValues = malloc(sizeof(OrtValue*) * numOutputs);
OrtValue** outputValues = allocarray(numOutputs, sizeof(OrtValue*));
if (outputValues == NULL) {
goto cleanup_input_values;
}
@ -512,11 +523,14 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep
// Release the java array copy of pointers to the tensors.
(*jniEnv)->ReleaseLongArrayElements(jniEnv, inputHandles, inputValueLongs, JNI_ABORT);
// Extract a C array of longs which are pointers to the output tensors.
jlong* outputHandleLongs = (*jniEnv)->GetLongArrayElements(jniEnv, outputHandlesArr, NULL);
// 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;
outputValues[i] = (OrtValue*)outputHandleLongs[i];
}
// Actually score the inputs.
@ -530,24 +544,29 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep
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);
// Create the output boolean array denoting if ORT owns the memory for each output.
// Java boolean arrays are initialized to false.
outputArray = (*jniEnv)->NewBooleanArray(jniEnv, safecast_int64_to_jsize(numOutputs));
jboolean* boolArr = (*jniEnv)->GetBooleanArrayElements(jniEnv, outputArray, NULL);
// Convert the output tensors into ONNXValues
for (int i = 0; i < numOutputs; i++) {
if (outputValues[i] != NULL) {
if (outputValues[i] != NULL && (*jniEnv)->GetObjectArrayElement(jniEnv, outputValuesArr, 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);
boolArr[i] = 1;
(*jniEnv)->SetObjectArrayElement(jniEnv, outputValuesArr, i, onnxValue);
}
}
// Write the output array back to Java.
(*jniEnv)->ReleaseBooleanArrayElements(jniEnv, outputArray, boolArr, 0);
// 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:
cleanup_output_values:
free(outputValues);
// Release the Java output strings
@ -561,15 +580,15 @@ JNIEXPORT jobjectArray JNICALL Java_ai_onnxruntime_OrtTrainingSession_evalStep
}
// Release the buffers
cleanup_input_values:
cleanup_input_values:
free((void*)inputValuePtrs);
cleanup_java_output_strings:
cleanup_java_output_strings:
free(javaOutputStrings);
cleanup_java_input_strings:
cleanup_java_input_strings:
free(javaInputStrings);
cleanup_output_names:
cleanup_output_names:
free((void*)outputNames);
cleanup_input_names:
cleanup_input_names:
free((void*)inputNames);
return outputArray;

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

@ -6,11 +6,14 @@ package ai.onnxruntime;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import ai.onnxruntime.OrtException.OrtErrorCode;
import ai.onnxruntime.OrtSession.Result;
import ai.onnxruntime.OrtSession.SessionOptions;
import ai.onnxruntime.OrtSession.SessionOptions.ExecutionMode;
@ -31,6 +34,8 @@ import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@ -71,7 +76,7 @@ public class InferenceTest {
@Test
public void testVersion() {
String version = env.getVersion();
Assertions.assertFalse(version.isEmpty());
assertFalse(version.isEmpty());
}
@Test
@ -749,6 +754,151 @@ public class InferenceTest {
}
}
@Test
public void testPinnedOutputs() throws OrtException {
String modelPath = TestHelpers.getResourcePath("/java-three-output-matmul.onnx").toString();
FloatBuffer outputABuf =
ByteBuffer.allocateDirect(4 * 4).order(ByteOrder.nativeOrder()).asFloatBuffer();
FloatBuffer outputBBuf =
ByteBuffer.allocateDirect(4 * 4).order(ByteOrder.nativeOrder()).asFloatBuffer();
FloatBuffer outputCBuf =
ByteBuffer.allocateDirect(4 * 4).order(ByteOrder.nativeOrder()).asFloatBuffer();
FloatBuffer tooSmallBuf =
ByteBuffer.allocateDirect(4 * 2).order(ByteOrder.nativeOrder()).asFloatBuffer();
FloatBuffer tooBigBuf =
ByteBuffer.allocateDirect(4 * 6).order(ByteOrder.nativeOrder()).asFloatBuffer();
FloatBuffer wrongShapeBuf =
ByteBuffer.allocateDirect(4 * 4).order(ByteOrder.nativeOrder()).asFloatBuffer();
LongBuffer wrongTypeBuf =
ByteBuffer.allocateDirect(8 * 4).order(ByteOrder.nativeOrder()).asLongBuffer();
try (SessionOptions options = new SessionOptions()) {
try (OrtSession session = env.createSession(modelPath, options);
OnnxTensor t = OnnxTensor.createTensor(env, new float[][] {{1, 2, 3, 4}});
OnnxTensor outputA = OnnxTensor.createTensor(env, outputABuf, new long[] {1, 4});
OnnxTensor outputB = OnnxTensor.createTensor(env, outputBBuf, new long[] {1, 4});
OnnxTensor outputC = OnnxTensor.createTensor(env, outputCBuf, new long[] {1, 4});
OnnxTensor tooSmall = OnnxTensor.createTensor(env, tooSmallBuf, new long[] {1, 2});
OnnxTensor tooBig = OnnxTensor.createTensor(env, tooBigBuf, new long[] {1, 6});
OnnxTensor wrongShape = OnnxTensor.createTensor(env, wrongShapeBuf, new long[] {2, 2});
OnnxTensor wrongType = OnnxTensor.createTensor(env, wrongTypeBuf, new long[] {1, 4})) {
Map<String, OnnxTensor> inputMap = Collections.singletonMap("input", t);
Set<String> requestedOutputs = new LinkedHashSet<>();
Map<String, OnnxTensor> pinnedOutputs = new LinkedHashMap<>();
// Test that all outputs can be pinned
pinnedOutputs.put("output-0", outputA);
pinnedOutputs.put("output-1", outputB);
pinnedOutputs.put("output-2", outputC);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
assertEquals(3, r.size());
assertSame(outputA, r.get(0));
assertSame(outputB, r.get(1));
assertSame(outputC, r.get(2));
assertFalse(r.isResultOwner(0));
assertFalse(r.isResultOwner(1));
assertFalse(r.isResultOwner(2));
// More tests
}
TestHelpers.zeroBuffer(outputABuf);
TestHelpers.zeroBuffer(outputBBuf);
TestHelpers.zeroBuffer(outputCBuf);
requestedOutputs.clear();
pinnedOutputs.clear();
// Test a single pinned output
pinnedOutputs.put("output-1", outputB);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
assertEquals(1, r.size());
assertSame(outputB, r.get(0));
assertSame(outputB, r.get("output-1").get());
assertFalse(r.isResultOwner(0));
// More tests
}
TestHelpers.zeroBuffer(outputABuf);
TestHelpers.zeroBuffer(outputBBuf);
TestHelpers.zeroBuffer(outputCBuf);
requestedOutputs.clear();
pinnedOutputs.clear();
// Test a mixture of pinned and generated outputs
requestedOutputs.add("output-0");
requestedOutputs.add("output-2");
pinnedOutputs.put("output-1", outputB);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
assertEquals(3, r.size());
// pinned outputs are first
assertSame(outputB, r.get(0));
assertSame(outputB, r.get("output-1").get());
// requested outputs are different
assertNotSame(outputA, r.get("output-0").get());
assertNotSame(outputC, r.get("output-2").get());
// check ownership.
assertFalse(r.isResultOwner(0));
assertTrue(r.isResultOwner(1));
assertTrue(r.isResultOwner(2));
// More tests
}
TestHelpers.zeroBuffer(outputABuf);
TestHelpers.zeroBuffer(outputBBuf);
TestHelpers.zeroBuffer(outputCBuf);
requestedOutputs.clear();
pinnedOutputs.clear();
// Test that overlapping names causes an error
requestedOutputs.add("output-1");
pinnedOutputs.put("output-1", outputB);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
fail("Should have thrown OrtException");
} catch (OrtException e) {
assertEquals(OrtErrorCode.ORT_JAVA_UNKNOWN, e.getCode());
}
requestedOutputs.clear();
pinnedOutputs.clear();
// Test that a tensor of the wrong type causes an error
pinnedOutputs.put("output-0", wrongType);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
fail("Should have thrown OrtException");
} catch (OrtException e) {
assertEquals(OrtErrorCode.ORT_INVALID_ARGUMENT, e.getCode());
}
requestedOutputs.clear();
pinnedOutputs.clear();
// Test that a tensor of the wrong shape (but right capacity) causes an error.
pinnedOutputs.put("output-1", wrongShape);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
fail("Should have thrown OrtException");
} catch (OrtException e) {
assertEquals(OrtErrorCode.ORT_INVALID_ARGUMENT, e.getCode());
}
requestedOutputs.clear();
pinnedOutputs.clear();
// Test that a tensor which is too small causes an error
pinnedOutputs.put("output-1", tooSmall);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
fail("Should have thrown OrtException");
} catch (OrtException e) {
assertEquals(OrtErrorCode.ORT_INVALID_ARGUMENT, e.getCode());
}
requestedOutputs.clear();
pinnedOutputs.clear();
// Test that a tensor which is too large causes an error
pinnedOutputs.put("output-1", tooBig);
try (OrtSession.Result r = session.run(inputMap, requestedOutputs, pinnedOutputs)) {
fail("Should have thrown OrtException");
} catch (OrtException e) {
assertEquals(OrtErrorCode.ORT_INVALID_ARGUMENT, e.getCode());
}
requestedOutputs.clear();
pinnedOutputs.clear();
}
}
}
private static File getTestModelsDir() throws IOException {
// get build directory, append downloaded models location
String cwd = System.getProperty("user.dir");

96
java/src/test/java/ai/onnxruntime/ModelGenerators.java
View file

@ -182,6 +182,102 @@ public final class ModelGenerators {
}
}
public void generateThreeOutputMatmul() throws IOException {
OnnxMl.GraphProto.Builder graph = OnnxMl.GraphProto.newBuilder();
graph.setName("ort-test-three-matmul");
// Add placeholders
OnnxMl.ValueInfoProto.Builder input = OnnxMl.ValueInfoProto.newBuilder();
input.setName("input");
OnnxMl.TypeProto inputType =
buildTensorTypeNode(
new long[] {-1, 4},
new String[] {"batch_size", null},
OnnxMl.TensorProto.DataType.FLOAT);
input.setType(inputType);
graph.addInput(input);
OnnxMl.ValueInfoProto.Builder outputA = OnnxMl.ValueInfoProto.newBuilder();
outputA.setName("output-0");
OnnxMl.TypeProto outputType =
buildTensorTypeNode(
new long[] {-1, 4},
new String[] {"batch_size", null},
OnnxMl.TensorProto.DataType.FLOAT);
outputA.setType(outputType);
graph.addOutput(outputA);
OnnxMl.ValueInfoProto.Builder outputB = OnnxMl.ValueInfoProto.newBuilder();
outputB.setName("output-1");
outputB.setType(outputType);
graph.addOutput(outputB);
OnnxMl.ValueInfoProto.Builder outputC = OnnxMl.ValueInfoProto.newBuilder();
outputC.setName("output-2");
outputC.setType(outputType);
graph.addOutput(outputC);
// Add initializers
OnnxMl.TensorProto.Builder tensor = OnnxMl.TensorProto.newBuilder();
tensor.addDims(4);
tensor.addDims(4);
Float[] floats =
new Float[] {1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, 10f, 11f, 12f, 13f, 14f, 15f, 16f};
tensor.addAllFloatData(Arrays.asList(floats));
tensor.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
tensor.setName("tensor");
graph.addInitializer(tensor);
OnnxMl.TensorProto.Builder addInit = OnnxMl.TensorProto.newBuilder();
addInit.addDims(4);
Float[] addFloats = new Float[] {1f, 2f, 3f, 4f};
addInit.addAllFloatData(Arrays.asList(addFloats));
addInit.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
addInit.setName("add-init");
graph.addInitializer(addInit);
// Add operations
OnnxMl.NodeProto.Builder matmul = OnnxMl.NodeProto.newBuilder();
matmul.setName("matmul-0");
matmul.setOpType("MatMul");
matmul.addInput("input");
matmul.addInput("tensor");
matmul.addOutput("matmul-output");
graph.addNode(matmul);
OnnxMl.NodeProto.Builder id = OnnxMl.NodeProto.newBuilder();
id.setName("id-1");
id.setOpType("Identity");
id.addInput("matmul-output");
id.addOutput("output-0");
graph.addNode(id);
OnnxMl.NodeProto.Builder add = OnnxMl.NodeProto.newBuilder();
add.setName("add-2");
add.setOpType("Add");
add.addInput("matmul-output");
add.addInput("add-init");
add.addOutput("output-1");
graph.addNode(add);
OnnxMl.NodeProto.Builder log = OnnxMl.NodeProto.newBuilder();
log.setName("log-3");
log.setOpType("Log");
log.addInput("matmul-output");
log.addOutput("output-2");
graph.addNode(log);
// Build model
OnnxMl.ModelProto.Builder model = OnnxMl.ModelProto.newBuilder();
model.setGraph(graph);
model.setDocString("ORT three output matmul test");
model.setModelVersion(0);
model.setIrVersion(8);
model.setDomain("ai.onnxruntime.test");
model.addOpsetImport(OnnxMl.OperatorSetIdProto.newBuilder().setVersion(18).build());
try (OutputStream os =
Files.newOutputStream(
Paths.get("src", "test", "resources", "java-three-output-matmul.onnx"))) {
model.build().writeTo(os);
}
}
private static void genCast(
String name,
OnnxMl.TensorProto.DataType inputDataType,

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

@ -262,6 +262,12 @@ public class TestHelpers {
return new File(TestHelpers.class.getResource(path).getFile()).toPath();
}
public static void zeroBuffer(FloatBuffer buf) {
for (int i = 0; i < buf.capacity(); i++) {
buf.put(i, 0.0f);
}
}
public static float[] loadTensorFromFile(Path filename) {
return loadTensorFromFile(filename, true);
}

12
java/src/test/java/ai/onnxruntime/TrainingTest.java
View file

@ -16,7 +16,6 @@ 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;
@ -69,8 +68,6 @@ public class TrainingTest {
}
}
// 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();
@ -99,14 +96,11 @@ public class TrainingTest {
ByteBuffer.allocateDirect(4 * expectedOutput.length)
.order(ByteOrder.nativeOrder())
.asFloatBuffer();
OnnxTensor outputTensor =
OnnxTensor.createTensor(env, output, new long[expectedOutput.length]);
OnnxTensor outputTensor = OnnxTensor.createTensor(env, output, new long[0]);
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);
try (OrtSession.Result r = trainingSession.trainStep(pinnedInputs, outputMap)) {
Assertions.assertEquals(expectedOutput[0], (float) outputTensor.getValue(), 1e-3f);
}
*/
} finally {
OnnxValue.close(outputMap);
OnnxValue.close(pinnedInputs);

BIN
java/src/test/resources/java-three-output-matmul.onnx Normal file
View file

Binary file not shown.