Adds ONNX export support to Tribuo's LinearSGDModels

Classification/SGD/pom.xml

@@ -60,6 +60,12 @@
             <type>test-jar</type>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>tribuo-onnx</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter</artifactId>

Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java

@@ -16,6 +16,7 @@
 
 package org.tribuo.classification.sgd.linear;
 
+import ai.onnx.proto.OnnxMl;
 import org.tribuo.Example;
 import org.tribuo.ImmutableFeatureMap;
 import org.tribuo.ImmutableOutputInfo;
@@ -26,10 +27,17 @@
 import org.tribuo.math.la.DenseMatrix;
 import org.tribuo.math.la.DenseVector;
 import org.tribuo.math.util.VectorNormalizer;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXExportable;
+import org.tribuo.onnx.ONNXOperators;
+import org.tribuo.onnx.ONNXShape;
+import org.tribuo.onnx.ONNXUtils;
 import org.tribuo.provenance.ModelProvenance;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -42,7 +50,7 @@
  * Proceedings of COMPSTAT, 2010.
  * </pre>
  */
-public class LinearSGDModel extends AbstractLinearSGDModel<Label> {
+public class LinearSGDModel extends AbstractLinearSGDModel<Label> implements ONNXExportable {
     private static final long serialVersionUID = 2L;
 
     private final VectorNormalizer normalizer;
@@ -101,6 +109,52 @@ protected String getDimensionName(int index) {
         return outputIDInfo.getOutput(index).getLabel();
     }
 
+    @Override
+    public OnnxMl.ModelProto exportONNXModel(String domain, long modelVersion) {
+        ONNXContext context = new ONNXContext();
+
+        // Build graph
+        OnnxMl.GraphProto graph = exportONNXGraph(context);
+
+        return innerExportONNXModel(graph,domain,modelVersion);
+    }
+
+    @Override
+    public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
+        OnnxMl.GraphProto.Builder graphBuilder = OnnxMl.GraphProto.newBuilder();
+
+        // Make inputs and outputs
+        OnnxMl.TypeProto inputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,featureIDMap.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto inputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(inputType).setName("input").build();
+        graphBuilder.addInput(inputValueProto);
+        OnnxMl.TypeProto outputType = ONNXUtils.buildTensorTypeNode(new ONNXShape(new long[]{-1,outputIDInfo.size()}, new String[]{"batch",null}), OnnxMl.TensorProto.DataType.FLOAT);
+        OnnxMl.ValueInfoProto outputValueProto = OnnxMl.ValueInfoProto.newBuilder().setType(outputType).setName("output").build();
+        graphBuilder.addOutput(outputValueProto);
+
+        // Add weights
+        OnnxMl.TensorProto weightInitializerProto = weightBuilder(context);
+        graphBuilder.addInitializer(weightInitializerProto);
+
+        // Add biases
+        OnnxMl.TensorProto biasInitializerProto = biasBuilder(context);
+        graphBuilder.addInitializer(biasInitializerProto);
+
+        // Make gemm
+        String[] gemmInputs = new String[]{inputValueProto.getName(),weightInitializerProto.getName(),biasInitializerProto.getName()};
+        OnnxMl.NodeProto gemm = ONNXOperators.GEMM.build(context,gemmInputs,new String[]{context.generateUniqueName("gemm_output")},Collections.emptyMap());
+        graphBuilder.addNode(gemm);
+
+        // Make output normalizer
+        List<OnnxMl.NodeProto> normalizerProtos = normalizer.exportNormalizer(context,gemm.getOutput(0),"output");
+        if (normalizerProtos.isEmpty()) {
+            throw new IllegalArgumentException("Normalizer " + normalizer.getClass() + " cannot be exported in ONNX models.");
+        } else {
+            graphBuilder.addAllNode(normalizerProtos);
+        }
+
+        return graphBuilder.build();
+    }
+
     private void readObject(java.io.ObjectInputStream in) throws IOException, ClassNotFoundException {
         in.defaultReadObject();
 

Classification/SGD/src/test/java/org/tribuo/classification/sgd/linear/TestSGDLinear.java

@@ -16,21 +16,32 @@
 
 package org.tribuo.classification.sgd.linear;
 
+import ai.onnxruntime.OrtEnvironment;
+import ai.onnxruntime.OrtException;
+import ai.onnxruntime.OrtSession;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.tribuo.Dataset;
 import org.tribuo.Example;
 import org.tribuo.Model;
+import org.tribuo.Prediction;
 import org.tribuo.Trainer;
+import org.tribuo.VariableIDInfo;
+import org.tribuo.VariableInfo;
 import org.tribuo.classification.Label;
+import org.tribuo.classification.LabelFactory;
 import org.tribuo.classification.evaluation.LabelEvaluation;
 import org.tribuo.classification.evaluation.LabelEvaluator;
 import org.tribuo.classification.example.LabelledDataGenerator;
 import org.tribuo.classification.sgd.objectives.Hinge;
+import org.tribuo.classification.sgd.objectives.LogMulticlass;
 import org.tribuo.common.sgd.AbstractLinearSGDTrainer;
 import org.tribuo.common.sgd.AbstractSGDTrainer;
 import org.tribuo.dataset.DatasetView;
+import org.tribuo.interop.onnx.DenseTransformer;
+import org.tribuo.interop.onnx.LabelTransformer;
+import org.tribuo.interop.onnx.ONNXExternalModel;
 import org.tribuo.math.optimisers.AdaGrad;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -39,6 +50,9 @@
 
 import java.io.IOException;
 import java.io.ObjectInputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
@@ -50,7 +64,7 @@
 
 public class TestSGDLinear {
 
-    private static final LinearSGDTrainer t = new LinearSGDTrainer(new Hinge(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED);
+    private static final LinearSGDTrainer t = new LinearSGDTrainer(new LogMulticlass(),new AdaGrad(0.1,0.1),5,1000, Trainer.DEFAULT_SEED);
 
     @BeforeAll
     public static void setup() {
@@ -61,8 +75,8 @@ public static void setup() {
         }
     }
 
-    public static Model<Label> testSGDLinear(Pair<Dataset<Label>,Dataset<Label>> p) {
-        Model<Label> m = t.train(p.getA());
+    public static LinearSGDModel testSGDLinear(Pair<Dataset<Label>,Dataset<Label>> p) {
+        LinearSGDModel m = (LinearSGDModel) t.train(p.getA());
         LabelEvaluator e = new LabelEvaluator();
         LabelEvaluation evaluation = e.evaluate(m,p.getB());
         Map<String, List<Pair<String,Double>>> features = m.getTopFeatures(3);
@@ -88,6 +102,56 @@ public void testSingleClassTraining() {
         assertEquals(1.0,evaluation.recall(new Label("Bar"))); // This is due to the random init of the classifier.
     }
 
+    @Test
+    public void testOnnxSerialization() throws IOException, OrtException {
+        Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();
+        LinearSGDModel model = (LinearSGDModel) t.train(p.getA());
+
+        // Write out model
+        Path onnxFile = Files.createTempFile("tribuo-sgd-test",".onnx");
+        model.saveONNXModel("org.tribuo.classification.sgd.linear.test",1,onnxFile);
+
+        // Prep mappings
+        Map<String, Integer> featureMapping = new HashMap<>();
+        for (VariableInfo f : model.getFeatureIDMap()){
+            VariableIDInfo id = (VariableIDInfo) f;
+            featureMapping.put(id.getName(),id.getID());
+        }
+        Map<Label, Integer> outputMapping = new HashMap<>();
+        for (Pair<Integer,Label> l : model.getOutputIDInfo()) {
+            outputMapping.put(l.getB(), l.getA());
+        }
+
+        // Load in via ORT
+        OrtEnvironment env = OrtEnvironment.getEnvironment();
+        env.close();
+        ONNXExternalModel<Label> onnxModel = ONNXExternalModel.createOnnxModel(new LabelFactory(),featureMapping,outputMapping,new DenseTransformer(),new LabelTransformer(),new OrtSession.SessionOptions(),onnxFile,"input");
+
+        // Generate predictions
+        List<Prediction<Label>> nativePredictions = model.predict(p.getB());
+        List<Prediction<Label>> onnxPredictions = onnxModel.predict(p.getB());
+
+        // Assert the predictions are identical
+        for (int i = 0; i < nativePredictions.size(); i++) {
+            Prediction<Label> tribuo = nativePredictions.get(i);
+            Prediction<Label> external = onnxPredictions.get(i);
+            assertEquals(tribuo.getOutput().getLabel(),external.getOutput().getLabel());
+            assertEquals(tribuo.getOutput().getScore(),external.getOutput().getScore(),1e-6);
+            for (Map.Entry<String,Label> l : tribuo.getOutputScores().entrySet()) {
+                Label other = external.getOutputScores().get(l.getKey());
+                if (other == null) {
+                    fail("Failed to find label " + l.getKey() + " in ORT prediction.");
+                } else {
+                    assertEquals(l.getValue().getScore(),other.getScore(),1e-6);
+                }
+            }
+        }
+
+        onnxModel.close();
+
+        onnxFile.toFile().delete();
+    }
+
     @Test
     public void testDenseData() {
         Pair<Dataset<Label>,Dataset<Label>> p = LabelledDataGenerator.denseTrainTest();

Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java

@@ -16,6 +16,8 @@
 
 package org.tribuo.common.sgd;
 
+import ai.onnx.proto.OnnxMl;
+import com.google.protobuf.ByteString;
 import com.oracle.labs.mlrg.olcut.util.Pair;
 import org.tribuo.Example;
 import org.tribuo.Excuse;
@@ -25,10 +27,16 @@
 import org.tribuo.Model;
 import org.tribuo.Output;
 import org.tribuo.Prediction;
+import org.tribuo.Tribuo;
 import org.tribuo.math.LinearParameters;
 import org.tribuo.math.la.DenseMatrix;
+import org.tribuo.onnx.ONNXContext;
+import org.tribuo.onnx.ONNXOperators;
 import org.tribuo.provenance.ModelProvenance;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.FloatBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -139,4 +147,71 @@ public Optional<Excuse<T>> getExcuse(Example<T> example) {
     public DenseMatrix getWeightsCopy() {
         return ((DenseMatrix)modelParameters.get()[0]).copy();
     }
+
+    /**
+     * Builds the ModelProto according to the standards for this model.
+     * @param graph The model graph.
+     * @param domain The model domain string.
+     * @param modelVersion The model version number.
+     * @return The ModelProto.
+     */
+    protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String domain, long modelVersion) {
+        // Build model
+        OnnxMl.ModelProto.Builder builder = OnnxMl.ModelProto.newBuilder();
+        builder.setGraph(graph);
+        builder.setDomain(domain);
+        builder.setProducerName("Tribuo");
+        builder.setProducerVersion(Tribuo.VERSION);
+        builder.setModelVersion(modelVersion);
+        builder.setDocString(toString());
+        builder.addOpsetImport(ONNXOperators.getOpsetProto());
+        builder.setIrVersion(6);
+        return builder.build();
+    }
+
+    /**
+     * Builds a TensorProto containing the model weights.
+     * @param context The naming context.
+     * @return The weight TensorProto.
+     */
+    protected OnnxMl.TensorProto weightBuilder(ONNXContext context) {
+        DenseMatrix weightMatrix = (DenseMatrix) modelParameters.get()[0];
+        OnnxMl.TensorProto.Builder weightBuilder = OnnxMl.TensorProto.newBuilder();
+        weightBuilder.setName(context.generateUniqueName("linear_sgd_weights"));
+        weightBuilder.addDims(featureIDMap.size());
+        weightBuilder.addDims(outputIDInfo.size());
+        weightBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
+        ByteBuffer buffer = ByteBuffer.allocate(featureIDMap.size() * outputIDInfo.size() * 4).order(ByteOrder.LITTLE_ENDIAN);
+        FloatBuffer floatBuffer = buffer.asFloatBuffer();
+        for (int j = 0; j < weightMatrix.getDimension2Size() - 1; j++) {
+            for (int i = 0; i < weightMatrix.getDimension1Size(); i++) {
+                floatBuffer.put((float) weightMatrix.get(i, j));
+            }
+        }
+        floatBuffer.rewind();
+        weightBuilder.setRawData(ByteString.copyFrom(buffer));
+        return weightBuilder.build();
+    }
+
+    /**
+     * Builds a TensorProto containing the model biases.
+     * @param context The naming context.
+     * @return The bias TensorProto.
+     */
+    protected OnnxMl.TensorProto biasBuilder(ONNXContext context) {
+        DenseMatrix weightMatrix = (DenseMatrix) modelParameters.get()[0];
+        OnnxMl.TensorProto.Builder biasBuilder = OnnxMl.TensorProto.newBuilder();
+        biasBuilder.setName(context.generateUniqueName("linear_sgd_biases"));
+        biasBuilder.addDims(outputIDInfo.size());
+        biasBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
+        ByteBuffer buffer = ByteBuffer.allocate(outputIDInfo.size()*4).order(ByteOrder.LITTLE_ENDIAN);
+        FloatBuffer floatBuffer = buffer.asFloatBuffer();
+        for (int i = 0; i < weightMatrix.getDimension1Size(); i++) {
+            floatBuffer.put((float)weightMatrix.get(i,weightMatrix.getDimension2Size()-1));
+        }
+        floatBuffer.rewind();
+        biasBuilder.setRawData(ByteString.copyFrom(buffer));
+        return biasBuilder.build();
+    }
+
 }

Core/pom.xml

@@ -37,6 +37,10 @@
             <groupId>com.oracle.labs.olcut</groupId>
             <artifactId>olcut-core</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.google.protobuf</groupId>
+            <artifactId>protobuf-java</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter</artifactId>

Core/src/main/java/ai/onnx/proto/NOTICE.md

@@ -0,0 +1,5 @@
+# OnnxMl.java
+
+This file was compiled from the onnx-ml.proto file in ONNX v1.9.0 using protoc 3.17.2.
+
+The onnx-ml.proto file is part of the ONNX project and licensed under Apache 2.0.