diff --git a/tools/python/transform/onnx_converter.py b/tools/python/transform/onnx_converter.py index c7d11ab72b2a5280a14c28bca9c0bce7c8cb8b1c..c1a7b0d680beb6836e70f5d20fe0cd79a089dbed 100644 --- a/tools/python/transform/onnx_converter.py +++ b/tools/python/transform/onnx_converter.py @@ -191,9 +191,9 @@ OnnxOpType = Enum('OnnxOpType', onnx_attr_translator = { "axis": lambda x: int(x), "axes": lambda x: [int(a) for a in x], - "dtype": lambda x: data_type.onnx2tf(x), + "dtype": lambda x: onnx_dtype(x), "keepdims": lambda x: bool(x), - "to": lambda x: data_type.onnx2tf(x), + "to": lambda x: onnx_dtype(x), } @@ -567,11 +567,7 @@ class OnnxConverter(base_converter.ConverterInterface): tensor.data_type = mace_pb2.DT_FLOAT tensor.float_data.extend( onnx_tensor.astype(np.float32).flat) - elif data_type == np.int32: - tensor.data_type = mace_pb2.DT_INT32 - tensor.int32_data.extend( - onnx_tensor.astype(np.int32).flat) - elif data_type == np.int64: + elif data_type == np.int64 or data_type == np.int32: tensor.data_type = mace_pb2.DT_INT32 tensor.int32_data.extend( onnx_tensor.astype(np.int32).flat) @@ -668,9 +664,9 @@ class OnnxConverter(base_converter.ConverterInterface): if 'to' in node.attrs: dtype = node.attrs['to'] - if dtype == TensorProto.FLOAT: + if dtype == np.float32 or dtype == np.float64: op.output_type.extend([self._option.data_type]) - elif dtype == TensorProto.INT: + elif dtype == np.int64 or dtype == np.int32: op.output_type.extend([mace_pb2.DT_INT32]) else: mace_check(False, "data type %s not supported" % dtype) @@ -959,7 +955,14 @@ class OnnxConverter(base_converter.ConverterInterface): if len(const_tensor.dims) == 0: value_arg = op.arg.add() value_arg.name = MaceKeyword.mace_scalar_input_str - value_arg.f = const_tensor.float_data[0] + if const_tensor.data_type == mace_pb2.DT_INT32: + value_arg.f = float(const_tensor.int32_data[0]) + elif const_tensor.data_type == mace_pb2.DT_FLOAT: + value_arg.f = const_tensor.float_data[0] + else: + mace_check(False, + "Does not support param's data type %s" + % const_tensor.data_type) value_index_arg = op.arg.add() value_index_arg.name = \ MaceKeyword.mace_scalar_input_index_str @@ -972,7 +975,14 @@ class OnnxConverter(base_converter.ConverterInterface): if len(const_tensor.dims) == 0: value_arg = op.arg.add() value_arg.name = MaceKeyword.mace_scalar_input_str - value_arg.f = const_tensor.float_data[0] + if const_tensor.data_type == mace_pb2.DT_INT32: + value_arg.f = float(const_tensor.int32_data[0]) + elif const_tensor.data_type == mace_pb2.DT_FLOAT: + value_arg.f = const_tensor.float_data[0] + else: + mace_check(False, + "Does not support param's data type %s" + % const_tensor.data_type) value_index_arg = op.arg.add() value_index_arg.name = \ MaceKeyword.mace_scalar_input_index_str diff --git a/tools/python/transform/shape_inference.py b/tools/python/transform/shape_inference.py index e035944d3231094f3dac5707dc775325da02e6be..ac97859bcdc82942c09d053e89f7ac2f2aa9663d 100644 --- a/tools/python/transform/shape_inference.py +++ b/tools/python/transform/shape_inference.py @@ -253,7 +253,7 @@ class ShapeInference(object): aspect_ratio = ConverterUtil.get_arg(op, MaceKeyword.mace_aspect_ratio_str).floats # noqa num_prior = len(aspect_ratio) * len(min_size) + len(max_size) - output_shape[2] = num_prior * input_h * input_w * 4 + output_shape[2] = int(num_prior * input_h * input_w * 4) self.add_output_shape(op, [output_shape]) def infer_shape_reshape(self, op): @@ -275,7 +275,7 @@ class ShapeInference(object): output_shape[i] = dim[i] product *= dim[i] if idx != -1: - output_shape[idx] = input_size / product + output_shape[idx] = int(input_size / product) self.add_output_shape(op, [output_shape]) else: output_shape = [] diff --git a/tools/python/transform/transformer.py b/tools/python/transform/transformer.py index 4c56b233f0f9395588d471168cdf4b350cc603c7..a080c12e904a1e21fd4340313773b976b785a53d 100644 --- a/tools/python/transform/transformer.py +++ b/tools/python/transform/transformer.py @@ -1440,15 +1440,15 @@ class Transformer(base_converter.ConverterInterface): arg.i = 1 elif arg.i == 3: arg.i = 2 - - producer = self._producer[op.input[0]] - input_shape = producer.output_shape[0].dims - if producer.type == MaceOp.FullyConnected.name and \ - len(input_shape) == 2: - axis_arg = ConverterUtil.get_arg( - op, MaceKeyword.mace_axis_str) - if axis_arg.i == 1: - axis_arg.i = 3 + if op.input[0] in self._producer: + producer = self._producer[op.input[0]] + input_shape = producer.output_shape[0].dims + if (producer.type == MaceOp.FullyConnected.name + and len(input_shape) == 2): + axis_arg = ConverterUtil.get_arg( + op, MaceKeyword.mace_axis_str) + if axis_arg.i == 1: + axis_arg.i = 3 elif op.type == MaceOp.Squeeze.name: for arg in op.arg: