add onnx export functionality

ibcn-cloudlet · Feb 15, 2018 · 68dbd26 · 68dbd26
1 parent 2ee6465
commit 68dbd26
Show file tree

Hide file tree

Showing 5 changed files with 622 additions and 2 deletions.
diff --git a/be.iminds.iot.dianne.api/src/be/iminds/iot/dianne/api/nn/module/dto/NeuralNetworkDTO.java b/be.iminds.iot.dianne.api/src/be/iminds/iot/dianne/api/nn/module/dto/NeuralNetworkDTO.java
@@ -72,4 +72,9 @@ public boolean equals(Object o){
 	public int hashCode(){
 		return name.hashCode();
 	}
+
+	@Override
+	public String toString() {
+		return name;
+	}
 }
diff --git a/be.iminds.iot.dianne.nn.test/models/.gitignore b/be.iminds.iot.dianne.nn.test/models/.gitignore
@@ -0,0 +1 @@
+*
diff --git a/be.iminds.iot.dianne.onnx/src/be/iminds/iot/dianne/onnx/DianneOnnx.java b/be.iminds.iot.dianne.onnx/src/be/iminds/iot/dianne/onnx/DianneOnnx.java
@@ -22,12 +22,16 @@
  *******************************************************************************/
 package be.iminds.iot.dianne.onnx;
 
+import java.util.Map;
+import java.util.UUID;
+
 import org.osgi.service.component.annotations.Component;
 import org.osgi.service.component.annotations.Reference;
 
 import be.iminds.iot.dianne.api.nn.module.dto.NeuralNetworkDTO;
 import be.iminds.iot.dianne.api.repository.DianneRepository;
 import be.iminds.iot.dianne.onnx.api.OnnxConverter;
+import be.iminds.iot.dianne.tensor.Tensor;
 
 @Component(immediate=true, 
 	property={"osgi.command.scope=dianne",
@@ -51,7 +55,15 @@ public NeuralNetworkDTO fromOnnx(String onnxFile) {
 	}
 
 	public void toOnnx(String onnxFile, String nnName, String... tag) {
-
+		try {
+			NeuralNetworkDTO nn = repository.loadNeuralNetwork(nnName);
+			Map<UUID, Tensor> parameters = repository.loadParameters(nnName, tag);
+			OnnxExporter exporter = new OnnxExporter(nn, parameters);
+			exporter.export(onnxFile);
+		} catch(Exception e) {
+			e.printStackTrace();
+			throw new RuntimeException("Failed to export "+nnName, e);
+		}
 	}
 
 }