def convert_ops(self, sess):
for tf_op in self._tf_graph.get_operations():
# print(__file__,sys._getframe().f_lineno,"tp_op=",tf_op)
mace_check(
tf_op.type in self._op_converters,
"Mace does not support tensorflow op type %s yet" % tf_op.type)
self._op_converters[tf_op.type](tf_op)
self.update_output_shapes(sess)
self.convert_tensors()
def convert_subpixel(self, tf_op):
op = self._mace_net_def.op.add()
index = tf_op.name.find('/')
op.name = tf_op.name[0:index] + '/Subpixel'
op.type = MaceOp.Subpixel.name
op.input.extend([self._subpixel_Ops["input"].inputs[0].name])
op.output.extend([
tf_output.name
for tf_output in self._subpixel_Ops["output"].outputs
])
inputDims = self.infer_tensor_shape_old(
self._subpixel_Ops["input"].inputs[0])
outputDims = self.infer_tensor_shape_old(
self._subpixel_Ops["output"].outputs[0])
scale = math.sqrt(inputDims[3] / outputDims[3])
subpiexl_scale_arg = op.arg.add()
subpiexl_scale_arg.name = MaceKeyword.mace_subpiexl_scale
subpiexl_scale_arg.i = int(scale)
for tf_output in self._subpixel_Ops["output"].outputs:
output_shape = op.output_shape.add()
self.infer_tensor_shape(output_shape, tf_output)
data_type_arg = op.arg.add()
data_type_arg.name = 'T'
try:
dtype = self._subpixel_Ops["input"].get_attr('T')
# print("dtype=",dtype)
if dtype == tf.int32:
data_type_arg.i = tensorrt_pb2.DT_INT32
elif dtype == tf.float32:
data_type_arg.i = self._option.data_type
else:
mace_check(False, "data type %s not supported" % dtype)
except ValueError:
try:
dtype = self._subpixel_Ops["input"].get_attr('SrcT')
if dtype == tf.int32 or dtype == tf.bool:
data_type_arg.i = tensorrt_pb2.DT_INT32
elif dtype == tf.float32:
data_type_arg.i = self._option.data_type
else:
mace_check(False, "data type %s not supported" % dtype)
except ValueError:
data_type_arg.i = self._option.data_type
ConverterUtil.add_data_format_arg(op, DataFormat.NHWC)
def convert_cast(self, tf_op):
op = self.convert_general_op(tf_op)
op.type = MaceOp.Cast.name
try:
dtype = tf_op.get_attr('DstT')
if dtype == tf.int32:
op.output_type.extend([tensorrt_pb2.DT_INT32])
elif dtype == tf.float32:
op.output_type.extend([self._option.data_type])
elif dtype == tf.uint8:
op.output_type.extend([tensorrt_pb2.DT_UINT8])
else:
mace_check(False, "data type %s not supported" % dtype)
except ValueError:
op.output_type.extend([self._option.data_type])
def convert_tensors(self):
for tf_op in self._tf_graph.get_operations():
if tf_op.type != TFOpType.Const.name:
continue
output_name = tf_op.outputs[0].name
if output_name not in self._skip_tensor:
tensor = self._mace_net_def.tensors.add()
tensor.name = tf_op.outputs[0].name
tf_tensor = tf_op.outputs[0].eval()
tensor.dims.extend(list(tf_tensor.shape))
tf_dt = tf_op.get_attr('dtype')
if tf_dt == tf.float32:
tensor.data_type = tensorrt_pb2.DT_FLOAT
tensor.float_data.extend(tf_tensor.astype(np.float32).flat)
elif tf_dt == tf.int32:
tensor.data_type = tensorrt_pb2.DT_INT32
tensor.int32_data.extend(tf_tensor.astype(np.int32).flat)
else:
mace_check(False,
"Not supported tensor type: %s" % tf_dt.name)
def convert_conv2d(self, tf_op):
op = self.convert_general_op(tf_op)
if tf_op.type == TFOpType.DepthwiseConv2dNative.name:
op.type = MaceOp.DepthwiseConv2d.name
elif tf_op.type == TFOpType.Conv2DBackpropInput.name:
op.type = MaceOp.Deconv2D.name
else:
op.type = MaceOp.Conv2D.name
padding_arg = op.arg.add()
padding_arg.name = MaceKeyword.mace_padding_str
padding_arg.i = self.padding_mode[tf_op.get_attr(tf_padding_str)].value
strides_arg = op.arg.add()
strides_arg.name = MaceKeyword.mace_strides_str
strides_arg.ints.extend(tf_op.get_attr(tf_strides_str)[1:3])
if op.type != MaceOp.Deconv2D.name:
dilation_arg = op.arg.add()
dilation_arg.name = MaceKeyword.mace_dilations_str
try:
dilation_val = tf_op.get_attr(tf_dilations_str)[1:3]
except ValueError:
dilation_val = [1, 1]
dilation_arg.ints.extend(dilation_val)
else:
mace_check(
len(tf_op.inputs) >= 3, "deconv should have (>=) 3 inputs.")
output_shape_arg = op.arg.add()
output_shape_arg.name = MaceKeyword.mace_output_shape_str
if tf_op.inputs[0].op.type == TFOpType.Const.name:
output_shape_value = \
tf_op.inputs[0].eval().astype(np.int32).flat
output_shape_arg.ints.extend(output_shape_value)
else:
output_shape_value = {}
output_shape_arg.ints.extend(output_shape_value)
del op.input[:]
op.input.extend([
tf_op.inputs[2].name, tf_op.inputs[1].name,
tf_op.inputs[0].name
])
def convert_general_op(self, tf_op):
op = self._mace_net_def.op.add()
op.name = tf_op.name
op.type = tf_op.type
op.input.extend([tf_input.name for tf_input in tf_op.inputs])
op.output.extend([tf_output.name for tf_output in tf_op.outputs])
for tf_output in tf_op.outputs:
output_shape = op.output_shape.add()
self.infer_tensor_shape(output_shape, tf_output)
data_type_arg = op.arg.add()
data_type_arg.name = 'T'
try:
dtype = tf_op.get_attr('T')
if dtype == tf.int32:
data_type_arg.i = tensorrt_pb2.DT_INT32
elif dtype == tf.float32:
data_type_arg.i = self._option.data_type
elif dtype == tf.uint8:
data_type_arg.i = tensorrt_pb2.DT_UINT8
else:
mace_check(False, "data type %s not supported" % dtype)
except ValueError:
try:
dtype = tf_op.get_attr('SrcT')
if dtype == tf.int32 or dtype == tf.bool:
data_type_arg.i = tensorrt_pb2.DT_INT32
elif dtype == tf.float32:
data_type_arg.i = self._option.data_type
elif dtype == tf.uint8:
data_type_arg.i = tensorrt_pb2.DT_UINT8
else:
mace_check(False, "data type %s not supported" % dtype)
except ValueError:
data_type_arg.i = self._option.data_type
ConverterUtil.add_data_format_arg(op, DataFormat.NHWC)
return op