def test_split(self):
"""Test split like operation's template."""
template = {
'var_0': {
'init': [
'self.{var_0} = P.Split(axis={axis}, output_num={output_num})'
],
'construct': ['opt_{var_0} = self.{var_0}({inputs})']
}
}
rewrite_data = {
'var_0': {
'operation': 'P.Split',
'variable_name': "split",
'output_type': 'array',
'inputs': ["x"],
'args': {
"axis": 1,
"output_num": 2
},
'weights': {},
'trainable_params': {}
}
}
fragment = Fragment(data_entity=rewrite_data,
code_template=template,
outputs=["opt_{var_0}"],
outputs_mapping=((0, 0), ))
code = fragment()
init = code[0]
construct = code[1]
self.assertEqual(init, ['self.split = P.Split(axis=1, output_num=2)'])
self.assertEqual(construct, ['opt_split = self.split(x)'])
self.assertEqual(fragment.get_outputs_by_idx(0, 1), 'opt_split[1]')
def test_transpose(self):
"""Test transpose like operation's template."""
template = {
'var_0': {
'init': ['self.{var_0} = P.Transpose()'],
'construct':
['opt_{var_0} = self.{var_0}({inputs}, (0, 3, 1, 2))']
}
}
rewrite_data = {
'var_0': {
'operation': 'P.Transpose',
'output_type': 'tensor',
'variable_name': "transpose",
'inputs': ["x"],
'args': {},
'weights': {},
'trainable_params': {}
}
}
fragment = Fragment(data_entity=rewrite_data,
code_template=template,
outputs=["opt_{var_0}"],
outputs_mapping=((0, 0), ))
code = fragment()
init = code[0]
construct = code[1]
self.assertEqual(init, ['self.transpose = P.Transpose()'])
self.assertEqual(construct,
['opt_transpose = self.transpose(x, (0, 3, 1, 2))'])
self.assertEqual(fragment.get_outputs_by_idx(0), "opt_transpose")
def test_matmul(self):
"""Test matmul like operation's template."""
template = {
'var_0': {
'init': [
'self.{var_0} = nn.MatMul()',
'self.{var_0}_w = Tensor(np.random.rand(*(2048, 1000)).astype(np.float32))'
],
'construct':
['opt_{var_0} = self.{var_0}({inputs},self.{var_0}_w)']
}
}
rewrite_data = {
'var_0': {
'operation': 'nn.MatMul',
'output_type': 'tensor',
'variable_name': "matmul",
'inputs': ["x"],
'args': {},
'weights': {},
'trainable_params': {}
},
'metadata': {
'source': 'probs/MatMul',
'operation': 'MatMul',
'scope': 'Model/MatMul',
'inputs': ['avg_pool/Mean:0', 'probs/MatMul/ReadVariableOp:0'],
'inputs_shape': (1, 2048),
'outputs': ['probs/MatMul:0'],
'outputs_shape': [1, 1000],
'precursor_nodes': ['avg_pool/Mean'],
'successor_nodes': ['probs/BiasAdd'],
'attributes': {}
}
}
fragment = Fragment(data_entity=rewrite_data,
code_template=template,
outputs=["opt_{var_0}"],
outputs_mapping=((0, 0), ))
code = fragment()
init = code[0]
construct = code[1]
self.assertEqual(init, [
'self.matmul = nn.MatMul()',
'self.matmul_w = Tensor(np.random.rand(*(2048, 1000)).astype(np.float32))'
])
self.assertEqual(construct,
['opt_matmul = self.matmul(x,self.matmul_w)'])
self.assertEqual(fragment.get_outputs_by_idx(0), "opt_matmul")
def test_biasadd(self):
"""Test biasadd like operation's template."""
template = {
'var_0': {
'init': [
'self.{var_0} = P.TensorAdd()',
'self.{var_0}_bias = Tensor(np.random.rand(*(1000,)).astype(np.float32))'
],
'construct':
['opt_{var_0} = self.{var_0}({inputs},self.{var_0}_bias)']
}
}
rewrite_data = {
'var_0': {
'operation': 'P.TensorAdd',
'output_type': 'tensor',
'variable_name': "add",
'inputs': ["x"],
'args': {},
'weights': {},
'trainable_params': {}
},
'metadata': {
'source': 'probs/BiasAdd',
'operation': 'Add',
'scope': 'Model/Add',
'inputs': ['probs/MatMul:0', 'probs/BiasAdd/ReadVariableOp:0'],
'inputs_shape': (1, 1000),
'outputs': ['probs/BiasAdd:0'],
'outputs_shape': [1, 1000],
'precursor_nodes': ['probs/MatMul'],
'successor_nodes': ['probs/Softmax'],
'attributes': {}
}
}
fragment = Fragment(data_entity=rewrite_data,
code_template=template,
outputs=["opt_{var_0}"],
outputs_mapping=((0, 0), ))
code = fragment()
init = code[0]
construct = code[1]
self.assertEqual(init, [
'self.add = P.TensorAdd()',
'self.add_bias = Tensor(np.random.rand(*(1000,)).astype(np.float32))'
])
self.assertEqual(construct, ['opt_add = self.add(x,self.add_bias)'])
self.assertEqual(fragment.get_outputs_by_idx(0), "opt_add")
def test_mapper(self):
"""Test code generation."""
onnx_info = {
"op_name":
"onnx::CumSum",
"attributes": {},
"weights": [
NodeWeight(weight_name="axis",
weight_value=np.int32(0),
weight_location=0)
]
}
template, exchange_msg, outputs_lists, outputs_mapping = ONNXToMindSporeMapper.convert(
onnx_info['op_name'], onnx_info['attributes'],
onnx_info['weights'])
exchange_msg['var_0']['variable_name'] = 'cumsum_op'
exchange_msg['var_0']['inputs'] = ['x']
fragment = Fragment(data_entity=exchange_msg,
code_template=template,
outputs=outputs_lists,
outputs_mapping=outputs_mapping)
code = fragment()
init_code = code[0]
construct_code = code[1]
self.assertEqual(init_code, [
"self.cumsum_op = P.CumSum(exclusive=False, reverse=False)",
"self.cumsum_op_axis = 0"
])
self.assertEqual(
construct_code,
["opt_cumsum_op = self.cumsum_op(x, self.cumsum_op_axis)"])
def batch_add_nodes(graph_obj, mapper) -> Generator:
"""
Add nodes to Generator in batch mode.
Args:
graph_obj (Graph): Graph obj.
mapper (Mapper): Mapper of third party framework and MindSpore.
"""
generator_inst = Generator()
external_inputs = graph_obj.user_provided_input_nodes
for node_name in graph_obj.nodes_in_topological_order:
node_inst = graph_obj.get_node(node_name)
node_input = graph_obj.get_input_shape(node_name)
node_output = graph_obj.get_output_shape(node_name)
if not node_input:
raise ValueError(
"Unable to get the node's inputs from Graph object.")
node_name_with_scope = _tf_model_node_name_reformat(
node_inst, node_name)
node_name = node_name_with_scope
node_inst.add_input_and_output_shape(node_input, node_output)
code_template, exchange_msg, outputs_lst, outputs_mapping = _convert_params(
node_inst, mapper, external_inputs)
outputs_mapping = NodeOutputManager(node_name,
output_mappings=outputs_mapping)
fragment = Fragment(data_entity=exchange_msg,
code_template=code_template,
outputs=outputs_lst,
outputs_mapping=outputs_mapping)
generator_inst.add_node(node_name,
node_instance=node_inst,
node_fragment=fragment)
return generator_inst
def bind_opt_var_names(self, fragment: Fragment):
"""Get the opt_var_name in return statement."""
for base_out in self._base_output_dict.values():
base_out.opt_var_name = fragment.get_outputs_by_idx(
base_out.idx_in_ms_provider)