def test_export_graph(self, model): onnx_graph = model.load().graph graph = OnnxImporter.import_graph(onnx_graph) exported_onnx_graph = OnnxExporter.export_graph(graph) imported_graph = OnnxImporter.import_graph(exported_onnx_graph) assert graph == imported_graph assert graph.opset == imported_graph.opset # ONNX exports the initializers in this model differently after importing - ONNX GS can't do much about this. if model.path != lstm_model().path: assert onnx_graph == exported_onnx_graph
def test_import_graph_tensor_map_preserved(self): model = identity_model() tensor_map = OrderedDict() graph = OnnxImporter.import_graph(model.load().graph, tensor_map=tensor_map) assert len(tensor_map) == 0 model.assert_equal(graph)
def test_import_graph_value_info(self): model = onnx.shape_inference.infer_shapes(identity_model().load()) graph = OnnxImporter.import_graph(model.graph) tensors = graph.tensors() assert all( [type(tensor) == Variable and tensor.dtype is not None and tensor.shape for tensor in tensors.values()] )
def import_onnx(onnx_model: "onnx.ModelProto") -> Graph: """ Import an onnx-graphsurgeon Graph from the provided ONNX model. Args: onnx_model (onnx.ModelProto): The ONNX model. Returns: Graph: A corresponding onnx-graphsurgeon Graph. """ from onnx_graphsurgeon.importers.onnx_importer import OnnxImporter return OnnxImporter.import_graph(onnx_model.graph, opset=OnnxImporter.get_opset(onnx_model))
def test_import_graph(self, model): graph = OnnxImporter.import_graph(model.load().graph) model.assert_equal(graph)
def test_import_graph_with_dim_param(self): model = dim_param_model() graph = OnnxImporter.import_graph(model.load().graph) model.assert_equal(graph)
def test_import_graph_with_initializer(self): model = lstm_model() graph = OnnxImporter.import_graph(model.load().graph) model.assert_equal(graph)
def test_export(self): with tempfile.NamedTemporaryFile() as f: onnx_model = gs.export_onnx(self.imported_graph) assert onnx_model assert OnnxImporter.import_graph( onnx_model.graph) == self.imported_graph
def setup_method(self): self.imported_graph = OnnxImporter.import_graph( identity_model().load().graph)