[Java] Adding model metadata support (#3573)

* java - adding deployment information to build.gradle.

* java - adding support for model metadata.
This commit is contained in:
Adam Pocock 2020-04-21 05:28:15 -04:00 committed by GitHub
parent 1c37d5e6ec
commit 3dd3f84116
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 417 additions and 10 deletions

2
.gitignore vendored
View file

@ -45,4 +45,4 @@ java/gradlew
java/gradlew.bat
java/gradle
java/.gradle
java/hs_*.log

View file

@ -1,6 +1,7 @@
plugins {
id 'java'
id 'jacoco'
id 'maven-publish'
id 'com.diffplug.gradle.spotless' version '3.26.0'
}
@ -13,8 +14,36 @@ allprojects {
java {
sourceCompatibility = JavaVersion.VERSION_1_8
targetCompatibility = JavaVersion.VERSION_1_8
withJavadocJar()
withSourcesJar()
}
project.group = "ai.onnxruntime"
version = rootProject.file('../VERSION_NUMBER').text.trim()
jar {
into("META-INF/maven/$project.group/$project.name") {
from { generatePomFileForMavenPublication }
rename ".*", "pom.xml"
}
}
// Add explicit sources jar with pom file.
task sourcesJar(type: Jar, dependsOn: classes) {
classifier = "sources"
from sourceSets.main.allSource
into("META-INF/maven/$project.group/$project.name") {
from { generatePomFileForMavenPublication }
rename ".*", "pom.xml"
}
}
// Add explicit javadoc jar with pom file
task javadocJar(type: Jar, dependsOn: javadoc) {
classifier = "javadoc"
from javadoc.destinationDir
into("META-INF/maven/$project.group/$project.name") {
from { generatePomFileForMavenPublication }
rename ".*", "pom.xml"
}
}
wrapper {
@ -40,8 +69,6 @@ def cmakeNativeLibDir = "${cmakeJavaDir}/native-lib"
def cmakeNativeJniDir = "${cmakeJavaDir}/native-jni"
def cmakeBuildOutputDir = "${cmakeJavaDir}/build"
version = rootProject.file('../VERSION_NUMBER').text.trim()
compileJava {
options.compilerArgs += ["-h", "${project.buildDir}/headers/"]
}
@ -76,6 +103,10 @@ if (cmakeBuildDir != null) {
}
task allJar(type: Jar) {
into("META-INF/maven/$project.group/$project.name") {
from { generatePomFileForMavenPublication }
rename ".*", "pom.xml"
}
classifier = 'all'
from sourceSets.main.output
from cmakeNativeJniDir
@ -96,7 +127,6 @@ if (cmakeBuildDir != null) {
cmakeBuild.dependsOn javadocJar
cmakeBuild.dependsOn javadoc
task cmakeCheck(type: Copy) {
from project.buildDir
include 'reports/**'
@ -116,6 +146,7 @@ test {
useJUnitPlatform()
testLogging {
events "passed", "skipped", "failed"
showStandardStreams = true
}
}
@ -126,3 +157,34 @@ jacocoTestReport {
html.destination file("${buildDir}/jacocoHtml")
}
}
publishing {
publications {
maven(MavenPublication) {
groupId = project.group
artifactId = project.name
from components.java
pom {
name = 'onnx-runtime'
description = 'ONNX Runtime is a performance-focused inference engine for ONNX (Open Neural Network Exchange) models.'
url = 'https://microsoft.github.io/onnxruntime/'
licenses {
license {
name = 'MIT License'
url = 'https://opensource.org/licenses/MIT'
}
}
organization {
name = 'Microsoft'
url = 'http://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'
}
}
}
}
}

View file

@ -0,0 +1,210 @@
package ai.onnxruntime;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
/**
* Contains the metadata associated with an ONNX model.
*
* <p>Unspecified default fields contain the empty string.
*
* <p>This class is a Java side copy of the native metadata, it does not access the native runtime.
*/
public final class OnnxModelMetadata {
private final String producerName;
private final String graphName;
private final String domain;
private final String description;
private final long version;
private final Map<String, String> customMetadata;
/**
* Constructed by an OrtSession in native code, nulls are replaced with the empty string or empty
* map as appropriate.
*
* @param producerName The model producer name.
* @param graphName The model graph name.
* @param domain The model domain name.
* @param description The model description.
* @param version The model version.
* @param customMetadataArray Any custom metadata associated with the model.
*/
OnnxModelMetadata(
String producerName,
String graphName,
String domain,
String description,
long version,
String[] customMetadataArray) {
this.producerName = producerName == null ? "" : producerName;
this.graphName = graphName == null ? "" : graphName;
this.domain = domain == null ? "" : domain;
this.description = description == null ? "" : description;
this.version = version;
if (customMetadataArray != null && customMetadataArray.length > 0) {
this.customMetadata = new HashMap<>();
if (customMetadataArray.length % 2 == 1) {
throw new IllegalStateException(
"Asked for keys and values, but received an odd number of elements.");
}
for (int i = 0; i < customMetadataArray.length; i += 2) {
customMetadata.put(customMetadataArray[i], customMetadataArray[i + 1]);
}
} else {
this.customMetadata = Collections.emptyMap();
}
}
/**
* Constructed by an OrtSession, nulls are replaced with the empty string or empty map as
* appropriate.
*
* @param producerName The model producer name.
* @param graphName The model graph name.
* @param domain The model domain name.
* @param description The model description.
* @param version The model version.
* @param customMetadata Any custom metadata associated with the model.
*/
OnnxModelMetadata(
String producerName,
String graphName,
String domain,
String description,
long version,
Map<String, String> customMetadata) {
this.producerName = producerName == null ? "" : producerName;
this.graphName = graphName == null ? "" : graphName;
this.domain = domain == null ? "" : domain;
this.description = description == null ? "" : description;
this.version = version;
this.customMetadata = customMetadata == null ? Collections.emptyMap() : customMetadata;
}
/**
* Copy constructor.
*
* @param other The metadata to copy.
*/
public OnnxModelMetadata(OnnxModelMetadata other) {
this.producerName = other.producerName;
this.graphName = other.graphName;
this.domain = other.domain;
this.description = other.description;
this.version = other.version;
this.customMetadata =
other.customMetadata.isEmpty()
? Collections.emptyMap()
: new HashMap<>(getCustomMetadata());
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
OnnxModelMetadata that = (OnnxModelMetadata) o;
return version == that.version
&& producerName.equals(that.producerName)
&& graphName.equals(that.graphName)
&& domain.equals(that.domain)
&& description.equals(that.description)
&& customMetadata.equals(that.customMetadata);
}
@Override
public int hashCode() {
return Objects.hash(producerName, graphName, domain, description, version, customMetadata);
}
/**
* Gets the producer name.
*
* @return The producer name.
*/
public String getProducerName() {
return producerName;
}
/**
* Gets the graph name.
*
* @return The graph name.
*/
public String getGraphName() {
return graphName;
}
/**
* Gets the domain.
*
* @return The domain.
*/
public String getDomain() {
return domain;
}
/**
* Gets the model description.
*
* @return The description.
*/
public String getDescription() {
return description;
}
/**
* Gets the model version.
*
* @return The model version.
*/
public long getVersion() {
return version;
}
/**
* Gets an unmodifiable reference to the complete custom metadata.
*
* @return The custom metadata.
*/
public Map<String, String> getCustomMetadata() {
return Collections.unmodifiableMap(customMetadata);
}
/**
* Returns Optional.of(value) if the custom metadata has a value for the supplied key, otherwise
* returns {@link Optional#empty}.
*
* @param key The custom metadata key.
* @return The custom metadata value if present.
*/
public Optional<String> getCustomMetadataValue(String key) {
return Optional.ofNullable(customMetadata.get(key));
}
@Override
public String toString() {
return "OnnxModelMetadata{"
+ "producerName='"
+ producerName
+ '\''
+ ", graphName='"
+ graphName
+ '\''
+ ", domain='"
+ domain
+ '\''
+ ", description='"
+ description
+ '\''
+ ", version="
+ version
+ ", customMetadata="
+ customMetadata
+ '}';
}
}

View file

@ -24,6 +24,8 @@ final class OnnxRuntime {
private static final int ORT_API_VERSION_1 = 1;
// Post 1.0 builds of the ORT API.
private static final int ORT_API_VERSION_2 = 2;
// Post 1.3 builds of the ORT API
private static final int ORT_API_VERSION_3 = 3;
/** The short name of the ONNX runtime shared library */
static final String ONNXRUNTIME_LIBRARY_NAME = "onnxruntime";
@ -50,7 +52,7 @@ final class OnnxRuntime {
try {
load(tempDirectory, ONNXRUNTIME_LIBRARY_NAME);
load(tempDirectory, ONNXRUNTIME_JNI_LIBRARY_NAME);
ortApiHandle = initialiseAPIBase(ORT_API_VERSION_2);
ortApiHandle = initialiseAPIBase(ORT_API_VERSION_3);
loaded = true;
} finally {
if (!isAndroid()) {

View file

@ -47,6 +47,8 @@ public class OrtSession implements AutoCloseable {
private final Set<String> outputNames;
private OnnxModelMetadata metadata;
private boolean closed = false;
/**
@ -201,7 +203,7 @@ public class OrtSession implements AutoCloseable {
/**
* Scores an input feed dict, returning the map of requested inferred outputs.
*
* <p>The outputs are sorted based on the supplied set traveral order.
* <p>The outputs are sorted based on the supplied set traversal order.
*
* @param inputs The inputs to score.
* @param requestedOutputs The requested outputs.
@ -263,6 +265,18 @@ public class OrtSession implements AutoCloseable {
}
}
/**
* Gets the metadata for the currently loaded model.
*
* @return The metadata.
*/
public OnnxModelMetadata getMetadata() throws OrtException {
if (metadata == null) {
metadata = constructMetadata(OnnxRuntime.ortApiHandle, nativeHandle, allocator.handle);
}
return metadata;
}
@Override
public String toString() {
return "OrtSession(numInputs=" + numInputs + ",numOutputs=" + numOutputs + ")";
@ -334,6 +348,18 @@ public class OrtSession implements AutoCloseable {
private native void closeSession(long apiHandle, long nativeHandle) throws OrtException;
/**
* Builds the {@link OnnxModelMetadata} for this session.
*
* @param ortApiHandle The api pointer.
* @param nativeHandle The native session pointer.
* @param allocatorHandle The OrtAllocator pointer.
* @return The metadata.
* @throws OrtException If the native runtime failed to access or allocate the metadata.
*/
private native OnnxModelMetadata constructMetadata(
long ortApiHandle, long nativeHandle, long allocatorHandle) throws OrtException;
/**
* Represents the options used to construct this session.
*

View file

@ -337,3 +337,94 @@ JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_releaseNamesHandle
}
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,names));
}
/*
* Class: ai_onnxruntime_OrtSession
* Method: constructMetadata
* Signature: (JJJ)Ljava/lang/String;
*/
JNIEXPORT jstring JNICALL Java_ai_onnxruntime_OrtSession_constructMetadata
(JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong nativeHandle, 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;
OrtAllocator* allocator = (OrtAllocator*) allocatorHandle;
// Setup
char* stringClassName = "java/lang/String";
jclass stringClazz = (*jniEnv)->FindClass(jniEnv, stringClassName);
char *metadataClassName = "ai/onnxruntime/OnnxModelMetadata";
jclass metadataClazz = (*jniEnv)->FindClass(jniEnv, metadataClassName);
//OnnxModelMetadata(String producerName, String graphName, String domain, String description, long version, String[] customMetadataArray)
jmethodID metadataConstructor = (*jniEnv)->GetMethodID(jniEnv, metadataClazz, "<init>",
"(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;J[Ljava/lang/String;)V");
// Get metadata
OrtModelMetadata* metadata;
checkOrtStatus(jniEnv,api,api->SessionGetModelMetadata((OrtSession*)nativeHandle,&metadata));
// Read out the producer name and convert it to a java.lang.String
char* charBuffer;
checkOrtStatus(jniEnv,api,api->ModelMetadataGetProducerName(metadata, allocator, &charBuffer));
jstring producerStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
// Read out the graph name and convert it to a java.lang.String
checkOrtStatus(jniEnv,api,api->ModelMetadataGetGraphName(metadata, allocator, &charBuffer));
jstring graphStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
// Read out the domain and convert it to a java.lang.String
checkOrtStatus(jniEnv,api,api->ModelMetadataGetDomain(metadata, allocator, &charBuffer));
jstring domainStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
// Read out the description and convert it to a java.lang.String
checkOrtStatus(jniEnv,api,api->ModelMetadataGetDescription(metadata, allocator, &charBuffer));
jstring descriptionStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
// Read out the version
int64_t version;
checkOrtStatus(jniEnv,api,api->ModelMetadataGetVersion(metadata, &version));
// Read out the keys, look up the values.
int64_t numKeys;
char** keys;
checkOrtStatus(jniEnv,api,api->ModelMetadataGetCustomMetadataMapKeys(metadata, allocator, &keys, &numKeys));
jobjectArray customArray = NULL;
if (numKeys > 0) {
customArray = (*jniEnv)->NewObjectArray(jniEnv,numKeys*2,stringClazz,NULL);
// Iterate key array to extract the values
for (int64_t i = 0; i < numKeys; i++) {
// Create a java.lang.String for the key
jstring keyJava = (*jniEnv)->NewStringUTF(jniEnv,keys[i]);
// Extract the value and convert it to a java.lang.String
checkOrtStatus(jniEnv,api,api->ModelMetadataLookupCustomMetadataMap(metadata,allocator,keys[i],&charBuffer));
jstring valueJava = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
// Write the key and value into the array
(*jniEnv)->SetObjectArrayElement(jniEnv, customArray, i*2, keyJava);
(*jniEnv)->SetObjectArrayElement(jniEnv, customArray, (i*2)+1, valueJava);
}
// Release key array
for (int64_t i = 0; i < numKeys; i++) {
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,keys[i]));
}
checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,keys));
} else {
customArray = (*jniEnv)->NewObjectArray(jniEnv,0,stringClazz,NULL);
}
// Invoke the metadata constructor
//OnnxModelMetadata(String producerName, String graphName, String domain, String description, long version, String[] customMetadataArray)
jobject metadataJava = (*jniEnv)->NewObject(jniEnv, metadataClazz, metadataConstructor, producerStr, graphStr, domainStr, descriptionStr, (jlong) version, customArray);
// Release the metadata
api->ReleaseModelMetadata(metadata);
return metadataJava;
}

View file

@ -41,6 +41,8 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.function.BiFunction;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;
@ -63,6 +65,7 @@ public class InferenceTest {
@Test
public void repeatedCloseTest() throws OrtException {
Logger.getLogger(OrtEnvironment.class.getName()).setLevel(Level.SEVERE);
OrtEnvironment env = OrtEnvironment.getEnvironment("repeatedCloseTest");
try (OrtEnvironment otherEnv = OrtEnvironment.getEnvironment()) {
assertFalse(otherEnv.isClosed());
@ -115,7 +118,9 @@ public class InferenceTest {
@Test
public void morePartialInputsTest() throws OrtException {
String modelPath = getResourcePath("/partial-inputs-test-2.onnx").toString();
try (OrtEnvironment env = OrtEnvironment.getEnvironment("partialInputs");
try (OrtEnvironment env =
OrtEnvironment.getEnvironment(
OrtEnvironment.LoggingLevel.ORT_LOGGING_LEVEL_FATAL, "partialInputs");
OrtSession.SessionOptions options = new SessionOptions();
OrtSession session = env.createSession(modelPath, options)) {
assertNotNull(session);
@ -200,7 +205,9 @@ public class InferenceTest {
@Test
public void partialInputsTest() throws OrtException {
String modelPath = getResourcePath("/partial-inputs-test.onnx").toString();
try (OrtEnvironment env = OrtEnvironment.getEnvironment("partialInputs");
try (OrtEnvironment env =
OrtEnvironment.getEnvironment(
OrtEnvironment.LoggingLevel.ORT_LOGGING_LEVEL_FATAL, "partialInputs");
OrtSession.SessionOptions options = new SessionOptions();
OrtSession session = env.createSession(modelPath, options)) {
assertNotNull(session);
@ -370,6 +377,15 @@ public class InferenceTest {
for (int i = 0; i < expectedOutputDimensions.length; i++) {
assertEquals(expectedOutputDimensions[i], outputInfo.shape[i]);
}
// Check the metadata can be extracted
OnnxModelMetadata metadata = session.getMetadata();
assertEquals("onnx-caffe2", metadata.getProducerName());
assertEquals("squeezenet_old", metadata.getGraphName());
assertEquals("", metadata.getDomain());
assertEquals("", metadata.getDescription());
assertEquals(0x7FFFFFFFFFFFFFFFL, metadata.getVersion());
assertTrue(metadata.getCustomMetadata().isEmpty());
}
}
}