Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor onnx math #187

Merged
merged 4 commits into from
Oct 27, 2021
Merged

Refactor onnx math #187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8a771f2
Refactoring methods that create `OnnxMl.TensorProto` instances into
JackSullivan Oct 14, 2021
846fd00
Fix `OnnxMathUtils.floatMatrixBuilder` to properly transpose non-squa…
JackSullivan Oct 22, 2021
d14f923
Updates based on comments
JackSullivan Oct 25, 2021
56db2fe
Merge branch 'main' into refactor-onnx-math-2
JackSullivan Oct 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions Classification/SGD/src/main/java/org/tribuo/classification/sgd/fm/FMClassificationModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
import org.tribuo.onnx.ONNXContext;
import org.tribuo.onnx.ONNXExportable;
import org.tribuo.onnx.ONNXShape;
import org.tribuo.onnx.ONNXUtils;
import org.tribuo.math.onnx.ONNXMathUtils;
import org.tribuo.provenance.ModelProvenance;

import java.util.LinkedHashMap;
Expand Down Expand Up @@ -116,10 +116,10 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
graphBuilder.setName("FMClassificationModel");

// 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.TypeProto inputType = ONNXMathUtils.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.TypeProto outputType = ONNXMathUtils.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);

Expand Down
7 changes: 3 additions & 4 deletions Classification/SGD/src/main/java/org/tribuo/classification/sgd/linear/LinearSGDModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,10 @@
import org.tribuo.onnx.ONNXExportable;
import org.tribuo.onnx.ONNXOperators;
import org.tribuo.onnx.ONNXShape;
import org.tribuo.onnx.ONNXUtils;
import org.tribuo.math.onnx.ONNXMathUtils;
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;
Expand Down Expand Up @@ -125,10 +124,10 @@ public OnnxMl.GraphProto exportONNXGraph(ONNXContext context) {
graphBuilder.setName("Classification-LinearSGDModel");

// 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.TypeProto inputType = ONNXMathUtils.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.TypeProto outputType = ONNXMathUtils.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);

Expand Down
75 changes: 6 additions & 69 deletions Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractFMModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
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;
Expand All @@ -27,15 +26,15 @@
import org.tribuo.Tribuo;
import org.tribuo.math.la.DenseMatrix;
import org.tribuo.math.la.DenseVector;
import org.tribuo.math.la.Matrix;
import org.tribuo.math.la.SGDVector;
import org.tribuo.math.la.Tensor;
import org.tribuo.math.onnx.ONNXMathUtils;
import org.tribuo.onnx.ONNXContext;
import org.tribuo.onnx.ONNXExportable;
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;
Expand Down Expand Up @@ -216,68 +215,6 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
return builder.build();
}

/**
* Builds a TensorProto containing the supplied DenseMatrix.
*
* @param context The ONNX context for naming.
* @param name The name for this tensor proto.
* @param matrix The matrix to store.
* @param transpose Should the matrix be transposed into the tensor?
* @return The matrix TensorProto.
*/
protected static OnnxMl.TensorProto matrixBuilder(ONNXContext context, String name, DenseMatrix matrix, boolean transpose) {
OnnxMl.TensorProto.Builder matrixBuilder = OnnxMl.TensorProto.newBuilder();
matrixBuilder.setName(context.generateUniqueName(name));
int dim1, dim2;
if (transpose) {
dim1 = matrix.getDimension2Size();
dim2 = matrix.getDimension1Size();
} else {
dim1 = matrix.getDimension1Size();
dim2 = matrix.getDimension2Size();
}
matrixBuilder.addDims(dim1);
matrixBuilder.addDims(dim2);
matrixBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
ByteBuffer buffer = ByteBuffer.allocate(dim1 * dim2 * 4).order(ByteOrder.LITTLE_ENDIAN);
FloatBuffer floatBuffer = buffer.asFloatBuffer();
for (int i = 0; i < dim1; i++) {
for (int j = 0; j < dim2; j++) {
if (transpose) {
floatBuffer.put((float) matrix.get(j, i));
} else {
floatBuffer.put((float) matrix.get(i, j));
}
}
}
floatBuffer.rewind();
matrixBuilder.setRawData(ByteString.copyFrom(buffer));
return matrixBuilder.build();
}

/**
* Builds a TensorProto containing the supplied dense vector.
*
* @param context The ONNX context for naming.
* @param name The name for this tensor proto.
* @param vector The vector to store.
* @return The vector TensorProto.
*/
protected static OnnxMl.TensorProto vectorBuilder(ONNXContext context, String name, DenseVector vector) {
OnnxMl.TensorProto.Builder vectorBuilder = OnnxMl.TensorProto.newBuilder();
vectorBuilder.setName(context.generateUniqueName(name));
vectorBuilder.addDims(vector.size());
vectorBuilder.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
ByteBuffer buffer = ByteBuffer.allocate(vector.size() * 4).order(ByteOrder.LITTLE_ENDIAN);
FloatBuffer floatBuffer = buffer.asFloatBuffer();
for (int i = 0; i < vector.size(); i++) {
floatBuffer.put((float) vector.get(i));
}
floatBuffer.rewind();
vectorBuilder.setRawData(ByteString.copyFrom(buffer));
return vectorBuilder.build();
}

/**
* Constructs the shared stem of the Factorization Machine, used by all output types.
* <p>
Expand All @@ -299,17 +236,17 @@ protected String generateONNXGraph(ONNXContext context, OnnxMl.GraphProto.Builde
graphBuilder.addInitializer(twoConst);

// Add weights
OnnxMl.TensorProto weightInitializerProto = matrixBuilder(context, "fm_linear_weights", (DenseMatrix) modelParams[1], true);
OnnxMl.TensorProto weightInitializerProto = ONNXMathUtils.floatMatrixBuilder(context, "fm_linear_weights", (Matrix) modelParams[1], true);
graphBuilder.addInitializer(weightInitializerProto);

// Add biases
OnnxMl.TensorProto biasInitializerProto = vectorBuilder(context, "fm_biases", (DenseVector) modelParams[0]);
OnnxMl.TensorProto biasInitializerProto = ONNXMathUtils.floatVectorBuilder(context, "fm_biases", (SGDVector) modelParams[0]);
graphBuilder.addInitializer(biasInitializerProto);

// Add embedding vectors
OnnxMl.TensorProto[] embeddingProtos = new OnnxMl.TensorProto[outputIDInfo.size()];
for (int i = 0; i < outputIDInfo.size(); i++) {
embeddingProtos[i] = matrixBuilder(context, "fm_embedding_" + i, (DenseMatrix) modelParams[i + 2], true);
embeddingProtos[i] = ONNXMathUtils.floatMatrixBuilder(context, "fm_embedding_" + i, (Matrix) modelParams[i + 2], true);
graphBuilder.addInitializer(embeddingProtos[i]);
}

Expand Down
51 changes: 19 additions & 32 deletions Common/SGD/src/main/java/org/tribuo/common/sgd/AbstractLinearSGDModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
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;
Expand All @@ -30,15 +29,16 @@
import org.tribuo.Tribuo;
import org.tribuo.math.LinearParameters;
import org.tribuo.math.la.DenseMatrix;
import org.tribuo.math.la.Matrix;
import org.tribuo.math.onnx.ONNXMathUtils;
import org.tribuo.onnx.ONNXContext;
import org.tribuo.onnx.ONNXExportable;
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.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
Expand Down Expand Up @@ -197,22 +197,15 @@ protected OnnxMl.ModelProto innerExportONNXModel(OnnxMl.GraphProto graph, String
* @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();
final Matrix weightMatrix = (Matrix) modelParameters.get()[0];
return ONNXMathUtils.floatTensorBuilder(context, "linear_sgd_weights", Arrays.asList(featureIDMap.size(), outputIDInfo.size()),
(FloatBuffer fb) -> {
for (int j = 0; j < weightMatrix.getDimension2Size() - 1; j++) {
for (int i = 0; i < weightMatrix.getDimension1Size(); i++) {
fb.put((float) weightMatrix.get(i, j));
}
}
});
}

/**
Expand All @@ -221,19 +214,13 @@ protected OnnxMl.TensorProto weightBuilder(ONNXContext 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();
Matrix weightMatrix = (Matrix) modelParameters.get()[0];
return ONNXMathUtils.floatTensorBuilder(context, "linear_sgd_biases", Collections.singletonList(outputIDInfo.size()),
(FloatBuffer fb) -> {
for (int i = 0; i < weightMatrix.getDimension1Size(); i++) {
fb.put((float)weightMatrix.get(i,weightMatrix.getDimension2Size()-1));
}
});
}

}
100 changes: 0 additions & 100 deletions Core/src/main/java/org/tribuo/onnx/ONNXUtils.java
View file
Open in desktop

This file was deleted.

2 changes: 0 additions & 2 deletions Core/src/main/java/org/tribuo/provenance/ModelProvenance.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import com.oracle.labs.mlrg.olcut.provenance.MapProvenance;
import com.oracle.labs.mlrg.olcut.provenance.ObjectProvenance;
import com.oracle.labs.mlrg.olcut.provenance.Provenance;
import com.oracle.labs.mlrg.olcut.provenance.ProvenanceException;
import com.oracle.labs.mlrg.olcut.provenance.primitives.DateTimeProvenance;
import com.oracle.labs.mlrg.olcut.provenance.primitives.StringProvenance;
import com.oracle.labs.mlrg.olcut.util.Pair;
Expand All @@ -32,7 +31,6 @@
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;

/**
* Contains provenance information for an instance of a {@link org.tribuo.Model}.
Expand Down
Loading