Add transforms to sas_onnx_parse
@@ -28,7 +28,10 @@ from onnx.shape_inference import infer_shapes
from dlpy.layers import (InputLayer, Conv2d, Pooling, Dense, OutputLayer,
BN, Concat, Res)
-
+from dlpy.model_conversion.onnx_graph import OnnxGraph, OnnxNode
+from dlpy.model_conversion.onnx_transforms import (ConstToInitializer,
+ InitReshape, InitUnsqueeze,
+ FuseMulAddBN)
# The supported ONNX ops that can be parsed by this module
_onnx_ops = ['Conv', 'MaxPool', 'AveragePool', 'GlobalAveragePool',
@@ -81,6 +84,18 @@ def onnx_to_sas(model, model_name=None, output_layer=None):
List of Layers
'''
+ # run transforms
+ graph_ = OnnxGraph.from_onnx(model.graph)
+ transforms = [
+ ConstToInitializer(),
+ InitReshape(),
+ InitUnsqueeze(),
+ FuseMulAddBN()
+ ]
+
+ for transform in transforms:
+ transform(graph_)
+ model = graph_.make_onnx()
# verify model ops are supported
for n in model.graph.node:
@@ -95,11 +110,6 @@ def onnx_to_sas(model, model_name=None, output_layer=None):
if model_name is None:
model_name = graph_def.name
- # create name if node doesn't have a name
- for idx, node in enumerate(graph_def.node):
- if not node.name:
- node.name = '{}_{}'.format(node.op_type, idx)
-
# nodes that correspond to sas deeplearn layers
sas_computation_nodes = onnx_filter_sas_layers(graph_def)
@@ -114,20 +124,6 @@ def onnx_to_sas(model, model_name=None, output_layer=None):
tensor_dict = dict(init_tensors)
- # check reshape ops following initialized tensor
- # if present, reshape and add to tensor_dict
- for node in graph_def.node:
- if node.op_type != 'Reshape':
- continue
- # data tensor, shape tensor
- data_name = node.input[0]
- shape_name = node.input[1]
- out_name = node.output[0]
- if data_name in initialized and shape_name in initialized:
- reshape_tensor = tensor_dict[shape_name].flatten().tolist()
- data_tensor = tensor_dict[data_name]
- tensor_dict[out_name] = data_tensor.reshape(tuple(reshape_tensor))
-
# determine SAS input layers from uninitialized input ValueInfo
uninitialized = [value_info for value_info in graph_def.input
if value_info.name not in initialized]
