diff --git a/mace/dsp/test/quantized_add_test.cc b/mace/dsp/test/quantized_add_test.cc index 6cf51bdb0bfb1704fa82cadcb4bc3237c23f97a1..f30d8424f68c1613064f3c7531b9685a41a0f215 100644 --- a/mace/dsp/test/quantized_add_test.cc +++ b/mace/dsp/test/quantized_add_test.cc @@ -17,6 +17,7 @@ static NetDef BuildNetDef() { input_op->set_type("INPUT"); input_op->set_node_id(0); input_op->set_padding(0); + input_op->add_out_max_byte_size(1000); // add op OperatorDef *add_op = net.add_op(); @@ -58,6 +59,10 @@ static NetDef BuildNetDef() { input_node_input->set_node_id(16); input_node_input->set_output_port(0); + add_op->add_out_max_byte_size(1000); + add_op->add_out_max_byte_size(1000); + add_op->add_out_max_byte_size(1000); + // output op OperatorDef *output_op = net.add_op(); output_op->set_name("__output__"); diff --git a/mace/dsp/test/quantized_maxpool_test.cc b/mace/dsp/test/quantized_maxpool_test.cc index f99b0b956edc70cfea8fd1f562b34dbc0a6c6d30..897d1ac432e955ebc3337424e708ad62d28ab2d9 100644 --- a/mace/dsp/test/quantized_maxpool_test.cc +++ b/mace/dsp/test/quantized_maxpool_test.cc @@ -24,6 +24,7 @@ static NetDef BuildNetDef(const vector &input_shape, input_op->set_type("INPUT"); input_op->set_node_id(0); input_op->set_padding(0); + input_op->add_out_max_byte_size(1000); // maxpool op OperatorDef *maxpool_op = net.add_op(); @@ -57,6 +58,9 @@ static NetDef BuildNetDef(const vector &input_shape, input_node_input = maxpool_op->add_node_input(); input_node_input->set_node_id(13); input_node_input->set_output_port(0); + maxpool_op->add_out_max_byte_size(1000); + maxpool_op->add_out_max_byte_size(1000); + maxpool_op->add_out_max_byte_size(1000); // output op OperatorDef *output_op = net.add_op(); diff --git a/mace/dsp/test/quantized_relu_test.cc b/mace/dsp/test/quantized_relu_test.cc index b1c8fbd099fe88acf4b2790966af91d3d45c0d37..685be71d9b51ab3a833579654da568552d310a0f 100644 --- a/mace/dsp/test/quantized_relu_test.cc +++ b/mace/dsp/test/quantized_relu_test.cc @@ -16,6 +16,7 @@ static NetDef BuildNetDef() { input_op->set_type("INPUT"); input_op->set_node_id(0); input_op->set_padding(0); + input_op->add_out_max_byte_size(1000); // relu op OperatorDef *relu_op = net.add_op(); @@ -38,6 +39,9 @@ static NetDef BuildNetDef() { input_node_input = relu_op->add_node_input(); input_node_input->set_node_id(11); input_node_input->set_output_port(0); + relu_op->add_out_max_byte_size(1000); + relu_op->add_out_max_byte_size(1000); + relu_op->add_out_max_byte_size(1000); // output op OperatorDef *output_op = net.add_op(); diff --git a/mace/dsp/test/quantized_resize_bilinear_test.cc b/mace/dsp/test/quantized_resize_bilinear_test.cc index 6ee9f404e8894d96cdc7b97d0001cb62f396c5d6..12a2f8d34b94aeb21c8d3507be4ab4b545c26c2e 100644 --- a/mace/dsp/test/quantized_resize_bilinear_test.cc +++ b/mace/dsp/test/quantized_resize_bilinear_test.cc @@ -18,6 +18,7 @@ static NetDef BuildNetDef() { input_op->set_type("INPUT"); input_op->set_node_id(0); input_op->set_padding(0); + input_op->add_out_max_byte_size(1200); // relu op OperatorDef *resize_bilinear_op = net.add_op(); @@ -45,6 +46,9 @@ static NetDef BuildNetDef() { input_node_input = resize_bilinear_op->add_node_input(); input_node_input->set_node_id(12); input_node_input->set_output_port(0); + resize_bilinear_op->add_out_max_byte_size(1200); + resize_bilinear_op->add_out_max_byte_size(1000); + resize_bilinear_op->add_out_max_byte_size(1000); // output op OperatorDef *output_op = net.add_op(); diff --git a/mace/dsp/test/supernode_test.cc b/mace/dsp/test/supernode_test.cc index 1262299d745684e2b2641aa037d6fbefeb095dfd..634795ecf5beb6adbb3f255666537897be960436 100644 --- a/mace/dsp/test/supernode_test.cc +++ b/mace/dsp/test/supernode_test.cc @@ -17,6 +17,7 @@ static NetDef BuildNetDef() { input_op->set_type("INPUT"); input_op->set_node_id(0); input_op->set_padding(0); + input_op->add_out_max_byte_size(1000); // add op OperatorDef *supernode_op = net.add_op(); @@ -76,6 +77,10 @@ static NetDef BuildNetDef() { input_node_input->set_node_id(20); input_node_input->set_output_port(0); + supernode_op->add_out_max_byte_size(1000); + supernode_op->add_out_max_byte_size(1000); + supernode_op->add_out_max_byte_size(1000); + // output op OperatorDef *output_op = net.add_op(); output_op->set_name("__output__"); diff --git a/mace/proto/mace.proto b/mace/proto/mace.proto index c1936693df3289b94e58b8aaec42121d04bb46b3..d3deb3e4bed4cf5a3755c0375d0efe93dd94c240 100644 --- a/mace/proto/mace.proto +++ b/mace/proto/mace.proto @@ -93,6 +93,7 @@ message OperatorDef { optional uint32 op_id = 101; optional uint32 padding = 102; repeated NodeInput node_input = 103; + repeated int32 out_max_byte_size = 104; // only support 32-bit len } // for memory optimization diff --git a/mace/python/tools/tf_dsp_converter_lib.py b/mace/python/tools/tf_dsp_converter_lib.py index c485a3669c86787ee45b29873bcf08e0f54a22c6..209173e90f8930d2fe7abbc767def55ac93e5e9a 100644 --- a/mace/python/tools/tf_dsp_converter_lib.py +++ b/mace/python/tools/tf_dsp_converter_lib.py @@ -132,6 +132,7 @@ def convert_ops(unresolved_ops, resolved_ops, net_def, output_node, dsp_ops): op_def.input.append(input_tensor.name) op_def.input.extend([t.name for t in s2b_op.inputs[1:]]) op_def.input.extend([min_tensor.name, max_tensor.name]) + op_def.out_max_byte_size.extend([max_elem_size(out) for out in quantize_op.outputs]) convert_op_outputs(op_def, quantize_op) elif has_padding_and_strides(first_op): op_def.padding = padding_mode[first_op.get_attr('padding')] @@ -143,6 +144,7 @@ def convert_ops(unresolved_ops, resolved_ops, net_def, output_node, dsp_ops): strides = first_op.get_attr('strides') strides_tensor = add_shape_const_node(net_def, first_op, strides, 'strides') op_def.input.extend([strides_tensor]) + op_def.out_max_byte_size.extend([max_elem_size(out) for out in first_op.outputs]) convert_op_outputs(op_def, first_op) elif is_node_flatten_reshape(first_op): op_def.type = 'Flatten' @@ -150,6 +152,7 @@ def convert_ops(unresolved_ops, resolved_ops, net_def, output_node, dsp_ops): convert_op_outputs(op_def, first_op) elif dsp_ops.has_op(first_op.type): op_def.input.extend([t.name for t in first_op.inputs]) + op_def.out_max_byte_size.extend([max_elem_size(out) for out in first_op.outputs]) convert_op_outputs(op_def, first_op) else: raise Exception('Unsupported op: ', first_op) @@ -201,9 +204,11 @@ def reverse_batch_to_space_and_biasadd(net_def): new_biasadd_op.input[0] = get_tensor_name_from_op(conv_requantize_op.name, 0) new_biasadd_op.input[2] = get_tensor_name_from_op(conv_requantize_op.name, 1) new_biasadd_op.input[3] = get_tensor_name_from_op(conv_requantize_op.name, 2) + new_biasadd_op.out_max_byte_size[0] = conv_requantize_op.out_max_byte_size[0] * 4 new_biasadd_requantize_op = mace_pb2.OperatorDef() new_biasadd_requantize_op.CopyFrom(biasadd_requantize_op) + new_biasadd_requantize_op.out_max_byte_size[0] = new_biasadd_op.out_max_byte_size[0] / 4 new_b2s_op = mace_pb2.OperatorDef() new_b2s_op.CopyFrom(b2s_op) @@ -320,6 +325,7 @@ def strip_input_quantize_and_output_dequantize(net_def, input_node, output_node) new_input_op.name = input_op.name new_input_op.type = input_op.type new_input_op.padding = input_op.padding + new_input_op.out_max_byte_size.extend([input_op.out_max_byte_size[0]/4, 4, 4]) new_ops.append(new_input_op) new_input_op.output_shape.extend([input_op.output_shape[0], minf_op.output_shape[0],