in tinynn/converter/operators/tflite/transformable.py [0:0]
def transform(self, graph_converter, mapping):
input_tensor = self.inputs[0]
weight_tensor = self.inputs[1]
output_tensor = self.outputs[0]
input_dim = len(input_tensor.shape)
weight_dim = len(weight_tensor.shape)
prev_ops = []
next_ops = []
if weight_dim == 3 or input_dim == 3:
self.stride.insert(0, 1)
self.padding.insert(0, 0)
self.dilation.insert(0, 1)
self.output_padding.insert(0, 0)
reshape_outputs = [
self.create_transform_tensor(
np.expand_dims(t.tensor, 2),
name=f'{self.outputs[0].name}_{t.name}_4d_input',
quantization=t.quantization,
)
for t in self.inputs[:2]
]
reshape_attrs = [self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in reshape_outputs]
reshape_ops = [
tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
for old, new, attr in zip(self.inputs[:2], reshape_outputs, reshape_attrs)
]
for op in reshape_ops:
op.extra_hints['direction'] = 'up'
if weight_dim == 3 and input_dim == 3:
prev_ops.extend(reshape_ops)
elif weight_dim == 3:
prev_ops.append(reshape_ops[1])
else:
prev_ops.append(reshape_ops[0])
conv_outputs = [
self.create_transform_tensor(
np.expand_dims(self.outputs[0].tensor, 2),
name=f'{self.outputs[0].name}_4d_output',
quantization=self.outputs[0].quantization,
)
]
conv_attrs = [self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in self.outputs[:1]]
conv_ops = [
tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
for old, new, attr in zip(conv_outputs, self.outputs[:1], conv_attrs)
]
for op in conv_ops:
op.extra_hints['direction'] = 'down'
next_ops.extend(conv_ops)
if weight_dim == 3 and input_dim == 3:
self.inputs = reshape_outputs + self.inputs[2:]
elif weight_dim == 3:
self.inputs = self.inputs[0:1] + reshape_outputs[1:2] + self.inputs[1:]
else:
self.inputs = reshape_outputs[0:1] + self.inputs[1:]
self.outputs = conv_outputs + self.outputs[1:]
weight_tensor = self.inputs[1]
elif weight_dim not in (4, 5):
assert False, "Only Conv[Transpose]1d/2d/3d is supported"
if output_tensor.shape[1] != weight_tensor.shape[1]:
warnings.warn(
'Group transposed conv is not supported if official tflite interpreter is used. If that is the case'
' for you, plese pass in `group_conv_rewrite=True`. If you want to run the model with TFLite micro,'
' then you may also need to pass in `tflite_micro_rewrite=True`'
)
if weight_dim in (3, 4):
assert all((x == 1 for x in self.dilation)), "Only dilation=1 is supported for conv_transpose2d"
if self.enable_mtk_ops:
conv_op = MTKTransposeConvOperator(
self.inputs[:2][::-1],
self.outputs,
depth_multiplier=1,
dilation_height_factor=self.dilation[0],
dilation_width_factor=self.dilation[1],
padding_type=tflite.Padding.VALID,
stride_height=self.stride[0],
stride_width=self.stride[1],
)
else:
conv_op = tfl_ops.TransposeConvOperator(
self.inputs[:2][::-1],
self.outputs,
strideH=self.stride[0],
strideW=self.stride[1],
padding=tflite.Padding.VALID,
fusedActivationFunction=self.fusedActivationFunction,
)
else:
conv_op = tfl_ops.Conv3dTransposeOperator(
self.inputs[:2][::-1],
self.outputs,
strideD=self.stride[0],
strideH=self.stride[1],
strideW=self.stride[2],
dilationDFactor=self.dilation[0],
dilationHFactor=self.dilation[1],
dilationWFactor=self.dilation[2],
padding=tflite.Padding.VALID,
fusedActivationFunction=self.fusedActivationFunction,
)
ops = self.wrap_ops_with_nhwc_nchw_transposes([conv_op], input_idx=1)
# Pad handling
output_shape = conv_op.outputs[0].shape
if sum(self.padding) > 0:
if weight_dim in (3, 4):
pad_h = self.padding[0]
pad_w = self.padding[1]
start = np.array([0, pad_h, pad_w, 0], dtype='int32')
pad_sizes = ((0, 0), (pad_h, pad_h), (pad_w, pad_w), (0, 0))
else:
pad_d = self.padding[0]
pad_h = self.padding[1]
pad_w = self.padding[2]
start = np.array([0, pad_d, pad_h, pad_w, 0], dtype='int32')
pad_sizes = ((0, 0), (pad_d, pad_d), (pad_h, pad_h), (pad_w, pad_w), (0, 0))
size = np.array(ops[1].outputs[0].shape, dtype='int32')
start_tensor = self.create_attr_tensor(start)
size_tensor = self.create_attr_tensor(size)
slice_out = ops[1].outputs[0]
pad_array = np.pad(self.outputs[0].tensor, pad_sizes)
slice_input = self.create_transform_tensor(pad_array, quantization=self.outputs[0].quantization)
ops[1].outputs[0] = slice_input
slice_op = tfl_ops.SliceOperator([slice_input, start_tensor, size_tensor], [slice_out])
output_shape = slice_input.shape
ops.insert(2, slice_op)
# Output shape handling
output_shape_tensor = self.create_attr_tensor(np.array(output_shape, dtype='int32'))
conv_op.inputs.insert(0, output_shape_tensor)
# Weight handling
weight = conv_op.inputs[1]
if weight_dim in (3, 4):
nchw2chwn_perm = np.array([1, 2, 3, 0], dtype='int32')
else:
nchw2chwn_perm = np.array([2, 3, 4, 1, 0], dtype='int32')
nchw2chwn_perm_tensor = self.create_attr_tensor(nchw2chwn_perm)
reordered_weight = self.create_transform_tensor(
np.transpose(weight.tensor, nchw2chwn_perm), quantization=weight.quantization
)
conv_op.inputs[1] = reordered_weight
reorder_op = tfl_ops.TransposeOperator([weight, nchw2chwn_perm_tensor], [reordered_weight])
ops.insert(1, reorder_op)
# Bias handling
if self.enable_mtk_ops or self.conv_transpose_with_bias:
kernel_num = output_tensor.shape[1]
if len(self.inputs) > 2 and self.inputs[2].shape[0] != kernel_num and self.inputs[2].shape[0] == 1:
if conv_op.inputs[-1].dtype == np.float32:
bias = torch.tensor([self.inputs[2][0]] * kernel_num, dtype='float32')
else:
bias = torch.tensor([self.inputs[2][0]] * kernel_num, dtype='int32')
conv_op.inputs.append(self.create_attr_tensor(bias))
else:
if len(self.inputs) == 2 or self.inputs[2] is None:
if conv_op.inputs[-1].dtype == np.dtype('float32'):
bias = np.zeros((kernel_num,), dtype='float32')
q_args = None
else:
bias = np.zeros((kernel_num,), dtype='int32')
else:
bias = self.inputs[2].tensor
q_args = None
if bias.dtype != np.dtype('float32'):
per_tensor = weight_tensor.quantization.dim is None
# Bias handling
if per_tensor:
bias_scale = input_tensor.quantization.scale * weight_tensor.quantization.scale
bias_zero_point = 0
bias_dim = None
else:
bias_scale = [input_tensor.quantization.scale * s for s in weight_tensor.quantization.scale]
bias_zero_point = [0] * len(bias_scale)
bias_dim = 0
q_args = QuantizationParameters(bias_scale, bias_zero_point, bias_dim)
conv_op.inputs.append(self.create_attr_tensor(bias, quantization=q_args))
else:
if len(self.inputs) > 2 and self.inputs[2] is not None:
bias_tensor = self.inputs[2]
add_out = ops[-2].outputs[0]
bias_transform = self.create_transform_tensor(
add_out.tensor.copy(), quantization=self.outputs[0].quantization
)
ops[-2].outputs[0] = bias_transform
ops.insert(len(ops) - 1, tfl_ops.AddOperator([bias_transform, bias_tensor], [add_out]))
ops = prev_ops + ops + next_ops
for op in ops:
graph_converter.add_operator(op)
graph_converter.try_restore_edges(mapping)