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 Java ONNX Interface #199

Merged
merged 8 commits into from
Dec 10, 2021
Merged

Refactor Java ONNX Interface #199

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
f008f22
Initial refactor of ONNXContext to manage relations between ONNX grap…
JackSullivan Dec 7, 2021
4657b74
enforcing traversal order for libSVM classifier
JackSullivan Dec 7, 2021
12e285a
Fixed issue in MultiLabelVotingCombiner where I unsqueezed input inst…
JackSullivan Dec 7, 2021
db9eedd
Factored `ONNXRef` and children out of `ONNXContext`.
JackSullivan Dec 7, 2021
b3bf665
Adding documentation and addressing comments from PR
JackSullivan Dec 8, 2021
070cdef
Merge branch 'main' of github.com:oracle/tribuo into refactor-onnx-graph
JackSullivan Dec 10, 2021
f662fe3
fixing documentation and merging changes from main
JackSullivan Dec 10, 2021
6ec8f97
added docs for `ONNXExportable` methods shadowed on `AbstractLinearSG…
JackSullivan Dec 10, 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
65 changes: 21 additions & 44 deletions ...on/Core/src/main/java/org/tribuo/classification/ensemble/FullyWeightedVotingCombiner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,18 @@

package org.tribuo.classification.ensemble;

import ai.onnx.proto.OnnxMl;
import com.oracle.labs.mlrg.olcut.provenance.ConfiguredObjectProvenance;
import com.oracle.labs.mlrg.olcut.provenance.impl.ConfiguredObjectProvenanceImpl;
import org.tribuo.Example;
import org.tribuo.ImmutableOutputInfo;
import org.tribuo.Prediction;
import org.tribuo.classification.Label;
import org.tribuo.ensemble.EnsembleCombiner;
import org.tribuo.onnx.ONNXContext;
import org.tribuo.onnx.ONNXNode;
import org.tribuo.onnx.ONNXOperators;
import org.tribuo.onnx.ONNXUtils;
import org.tribuo.onnx.ONNXRef;
import org.tribuo.onnx.ONNXInitializer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
Expand Down Expand Up @@ -134,68 +133,46 @@ public ConfiguredObjectProvenance getProvenance() {
}

/**
* Exports this voting combiner as a list of ONNX NodeProtos.
* Exports this voting combiner to ONNX.
* <p>
* The input should be a 3-tensor [batch_size, num_outputs, num_ensemble_members].
* @param context The ONNX context object for name generation.
* @param input The name of the input tensor to combine.
* @param output The name of the voting output.
* @return A list of node protos representing the voting operation.
* @param input the node to be ensembled according to this implementation.
* @return The leaf node of the voting operation.
*/
@Override
public List<OnnxMl.NodeProto> exportCombiner(ONNXContext context, String input, String output) {
List<OnnxMl.NodeProto> nodes = new ArrayList<>();

public ONNXNode exportCombiner(ONNXNode input) {
// Take the mean over the maxed predictions
Map<String,Object> attributes = new HashMap<>();
attributes.put("axes",new int[]{2});
attributes.put("keepdims",0);
OnnxMl.NodeProto mean = ONNXOperators.REDUCE_MEAN.build(context,input,output,attributes);
nodes.add(mean);

return nodes;
return input.apply(ONNXOperators.REDUCE_MEAN, attributes);
}

/**
* Exports this ensemble combiner as a list of ONNX NodeProtos.
* Exports this voting combiner to ONNX.
* <p>
* The input should be a 3-tensor [batch_size, num_outputs, num_ensemble_members].
* @param context The ONNX context object for name generation.
* @param input The name of the input tensor to combine.
* @param output The name of the voting output.
* @param weight The name of the combination weight initializer.
* @return A list of node protos representing the voting operation.
* @param input the node to be ensembled according to this implementation.
* @param weight The node of weights for ensembling.
* @return The leaf node of the voting operation.
*/
@Override
public List<OnnxMl.NodeProto> exportCombiner(ONNXContext context, String input, String output, String weight) {
List<OnnxMl.NodeProto> nodes = new ArrayList<>();

public <T extends ONNXRef<?>> ONNXNode exportCombiner(ONNXNode input, T weight) {
Craigacp marked this conversation as resolved.
Show resolved Hide resolved
// Unsqueeze the weights to make sure they broadcast how I want them too.
// Now the size is [1, 1, num_members].
OnnxMl.TensorProto unsqueezeAxes = ONNXUtils.arrayBuilder(context,"unsqueeze_ensemble_output",new long[]{0,1});
context.addInitializer(unsqueezeAxes);
OnnxMl.NodeProto unsqueeze = ONNXOperators.UNSQUEEZE.build(context,new String[]{weight,unsqueezeAxes.getName()},context.generateUniqueName("unsqueezed_weights"));
nodes.add(unsqueeze);
ONNXInitializer unsqueezeAxes = input.onnxContext().array("unsqueeze_ensemble_output", new long[]{0, 1});

ONNXNode unsqueezed = weight.apply(ONNXOperators.UNSQUEEZE, unsqueezeAxes);

// Multiply the input by the weights.
OnnxMl.NodeProto mulByWeights = ONNXOperators.MUL.build(context,new String[]{input,unsqueeze.getOutput(0)},context.generateUniqueName("mul_predictions_by_weights"));
nodes.add(mulByWeights);
ONNXNode mulByWeights = input.apply(ONNXOperators.MUL, unsqueezed);

// Sum the weights
OnnxMl.NodeProto weightSum = ONNXOperators.REDUCE_SUM.build(context,weight,context.generateUniqueName("ensemble_weight_sum"));
nodes.add(weightSum);
ONNXNode weightSum = weight.apply(ONNXOperators.REDUCE_SUM);

// Take the weighted mean over the outputs
OnnxMl.TensorProto sumAxes = ONNXUtils.arrayBuilder(context,"sum_across_ensemble_axes",new long[]{2});
context.addInitializer(sumAxes);
OnnxMl.NodeProto sumAcrossMembers = ONNXOperators.REDUCE_SUM.build(context,
new String[]{mulByWeights.getOutput(0),sumAxes.getName()},
context.generateUniqueName("sum_across_ensemble"),
Collections.singletonMap("keepdims",0));
nodes.add(sumAcrossMembers);
OnnxMl.NodeProto divideByWeightSum = ONNXOperators.DIV.build(context,new String[]{sumAcrossMembers.getOutput(0),weightSum.getOutput(0)},output);
nodes.add(divideByWeightSum);

return nodes;
ONNXInitializer sumAxes = input.onnxContext().array("sum_across_ensemble_axes", new long[]{2});
return mulByWeights.apply(ONNXOperators.REDUCE_SUM, sumAxes, Collections.singletonMap("keepdims", 0))
.apply(ONNXOperators.DIV, weightSum);
}
}
73 changes: 23 additions & 50 deletions Classification/Core/src/main/java/org/tribuo/classification/ensemble/VotingCombiner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,18 @@

package org.tribuo.classification.ensemble;

import ai.onnx.proto.OnnxMl;
import com.oracle.labs.mlrg.olcut.provenance.ConfiguredObjectProvenance;
import com.oracle.labs.mlrg.olcut.provenance.impl.ConfiguredObjectProvenanceImpl;
import org.tribuo.Example;
import org.tribuo.ImmutableOutputInfo;
import org.tribuo.Prediction;
import org.tribuo.classification.Label;
import org.tribuo.ensemble.EnsembleCombiner;
import org.tribuo.onnx.ONNXContext;
import org.tribuo.onnx.ONNXNode;
import org.tribuo.onnx.ONNXOperators;
import org.tribuo.onnx.ONNXUtils;
import org.tribuo.onnx.ONNXRef;
import org.tribuo.onnx.ONNXInitializer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
Expand Down Expand Up @@ -126,76 +125,50 @@ public ConfiguredObjectProvenance getProvenance() {
}

/**
* Exports this voting combiner as a list of ONNX NodeProtos.
* Exports this voting combiner to ONNX.
* <p>
* The input should be a 3-tensor [batch_size, num_outputs, num_ensemble_members].
* @param context The ONNX context object for name generation.
* @param input The name of the input tensor to combine.
* @param output The name of the voting output.
* @return A list of node protos representing the voting operation.
* @param input The input tensor to combine.
* @return the final node proto representing the voting operation.
*/
@Override
public List<OnnxMl.NodeProto> exportCombiner(ONNXContext context, String input, String output) {
List<OnnxMl.NodeProto> nodes = new ArrayList<>();

public ONNXNode exportCombiner(ONNXNode input) {
// Hardmax!
OnnxMl.NodeProto hardMaxEnsemble = ONNXOperators.HARDMAX.build(context,input,context.generateUniqueName("hardmax_predictions"),Collections.singletonMap("axis",1));
nodes.add(hardMaxEnsemble);

// Take the mean over the maxed predictions
Map<String,Object> attributes = new HashMap<>();
attributes.put("axes",new int[]{2});
attributes.put("keepdims",0);
OnnxMl.NodeProto mean = ONNXOperators.REDUCE_MEAN.build(context,hardMaxEnsemble.getOutput(0),output,attributes);
nodes.add(mean);

return nodes;
return input.apply(ONNXOperators.HARDMAX, Collections.singletonMap("axis", 1))
.apply(ONNXOperators.REDUCE_MEAN, attributes);
}

/**
* Exports this voting combiner as a list of ONNX NodeProtos.
* Exports this voting combiner to ONNX
* <p>
* The input should be a 3-tensor [batch_size, num_outputs, num_ensemble_members].
* @param context The ONNX context object for name generation.
* @param input The name of the input tensor to combine.
* @param output The name of the voting output.
* @param weight The name of the combination weight initializer.
* @return A list of node protos representing the voting operation.
* @param input The input tensor to combine.
* @param weight The combination weight node.
* @return the final node proto representing the voting operation.
*/
@Override
public List<OnnxMl.NodeProto> exportCombiner(ONNXContext context, String input, String output, String weight) {
List<OnnxMl.NodeProto> nodes = new ArrayList<>();

public <T extends ONNXRef<?>> ONNXNode exportCombiner(ONNXNode input, T weight) {
// Unsqueeze the weights to make sure they broadcast how I want them too.
// Now the size is [1, 1, num_members].
OnnxMl.TensorProto unsqueezeAxes = ONNXUtils.arrayBuilder(context,"unsqueeze_ensemble_output",new long[]{0,1});
context.addInitializer(unsqueezeAxes);
OnnxMl.NodeProto unsqueeze = ONNXOperators.UNSQUEEZE.build(context,new String[]{weight,unsqueezeAxes.getName()},context.generateUniqueName("unsqueezed_weights"));
nodes.add(unsqueeze);
ONNXInitializer unsqueezeAxes = input.onnxContext().array("unsqueeze_ensemble_output", new long[]{0, 1});
ONNXInitializer sumAxes = input.onnxContext().array("sum_across_ensemble_axes", new long[]{2});

// Hardmax!
OnnxMl.NodeProto hardMaxEnsemble = ONNXOperators.HARDMAX.build(context,input,context.generateUniqueName("hardmax_predictions"),Collections.singletonMap("axis",1));
nodes.add(hardMaxEnsemble);
ONNXNode unsqueezed = weight.apply(ONNXOperators.UNSQUEEZE, unsqueezeAxes);

// Hardmax!
// Multiply the input by the weights.
OnnxMl.NodeProto mulByWeights = ONNXOperators.MUL.build(context,new String[]{hardMaxEnsemble.getOutput(0),unsqueeze.getOutput(0)},context.generateUniqueName("mul_predictions_by_weights"));
nodes.add(mulByWeights);
ONNXNode mulByWeights = input.apply(ONNXOperators.HARDMAX, Collections.singletonMap("axis", 1))
.apply(ONNXOperators.MUL, unsqueezed);

// Sum the weights
OnnxMl.NodeProto weightSum = ONNXOperators.REDUCE_SUM.build(context,weight,context.generateUniqueName("ensemble_weight_sum"));
nodes.add(weightSum);
ONNXNode weightSum = weight.apply(ONNXOperators.REDUCE_SUM);

// Take the weighted mean over the outputs
OnnxMl.TensorProto sumAxes = ONNXUtils.arrayBuilder(context,"sum_across_ensemble_axes",new long[]{2});
context.addInitializer(sumAxes);
OnnxMl.NodeProto sumAcrossMembers = ONNXOperators.REDUCE_SUM.build(context,
new String[]{mulByWeights.getOutput(0),sumAxes.getName()},
context.generateUniqueName("sum_across_ensemble"),
Collections.singletonMap("keepdims",0));
nodes.add(sumAcrossMembers);
OnnxMl.NodeProto divideByWeightSum = ONNXOperators.DIV.build(context,new String[]{sumAcrossMembers.getOutput(0),weightSum.getOutput(0)},output);
nodes.add(divideByWeightSum);

return nodes;
return mulByWeights.apply(ONNXOperators.REDUCE_SUM, sumAxes, Collections.singletonMap("keepdims", 0))
.apply(ONNXOperators.DIV, weightSum);
}
}
Loading