Roshrini commented on a change in pull request #10605: [MXNET-310] [ONNX-MXNet]
API to import ONNX models into Gluon.
URL: https://github.com/apache/incubator-mxnet/pull/10605#discussion_r191900033
##########
File path: python/mxnet/contrib/onnx/_import/import_onnx.py
##########
@@ -155,6 +157,40 @@ def get_graph_metadata(self, graph):
}
return metadata
+ def graph_to_gluon(self, graph, context):
+ """Construct SymbolBlock from onnx graph.
+
+ Parameters
+ ----------
+ graph : onnx protobuf object
+ The loaded onnx graph
+ context : str
+ context for mxnet module object. Should be 'CPU' or 'GPU'
+
+ Returns
+ -------
+ sym_block :gluon.nn.SymbolBlock
+ The returned gluon SymbolBlock
+ """
+ sym, arg_params, aux_params = self.from_onnx(graph)
+ metadata = self.get_graph_metadata(graph)
+ data_names = [input_tensor[0] for input_tensor in
metadata['input_tensor_data']]
+ data_inputs = [symbol.var(data_name) for data_name in data_names]
+
+ ctx = gpu() if context == 'GPU' else cpu()
Review comment:
from .... import cpu, gpu
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
