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]
input_dim = len(input_tensor.shape)
weight_dim = len(weight_tensor.shape)
prev_ops = []
next_ops = []
if weight_dim == 3 or input_dim == 3:
reshape_input_size = 1
reshape_output_size = 1
if weight_dim == 3:
self.stride.insert(0, 1)
self.padding.insert(0, 0)
self.dilation.insert(0, 1)
self.output_padding.insert(0, 0)
reshape_input_size = 2
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[:reshape_input_size]
]
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[:reshape_input_size], reshape_outputs, reshape_attrs)
]
for op in reshape_ops:
op.extra_hints['direction'] = 'up'
prev_ops.extend(reshape_ops)
conv_outputs = [
self.create_transform_tensor(
np.expand_dims(self.outputs[i].tensor, 2),
name=f'{self.outputs[i].name}_4d_output',
quantization=self.outputs[i].quantization,
)
for i in range(reshape_output_size)
]
conv_attrs = [
self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in self.outputs[:reshape_output_size]
]
conv_ops = [
tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
for old, new, attr in zip(conv_outputs, self.outputs[:reshape_output_size], conv_attrs)
]
for op in conv_ops:
op.extra_hints['direction'] = 'down'
next_ops.extend(conv_ops)
self.inputs = reshape_outputs + self.inputs[reshape_input_size:]
self.outputs = conv_outputs + self.outputs[reshape_output_size:]
weight_tensor = self.inputs[1]
elif weight_dim not in (4, 5):
assert False, "Only Conv[Transpose]1d/2d/3d is supported"
if weight_tensor.shape[1] == 1 and weight_tensor.shape[0] == self.groups:
if weight_dim in (3, 4):
conv_op = tfl_ops.DepthwiseConv2dOperator(
self.inputs,
self.outputs,
strideH=self.stride[0],
strideW=self.stride[1],
depthMultiplier=1,
dilationHFactor=self.dilation[0],
dilationWFactor=self.dilation[1],
fusedActivationFunction=self.fusedActivationFunction,
padding=tflite.Padding.VALID,
)
else:
assert False, "Only DepthwiseConv1d/2d is supported"
else:
if input_tensor.shape[1] != weight_tensor.shape[1]:
warnings.warn(
'Group 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):
conv_op = tfl_ops.Conv2dOperator(
self.inputs,
self.outputs,
strideH=self.stride[0],
strideW=self.stride[1],
dilationHFactor=self.dilation[0],
dilationWFactor=self.dilation[1],
fusedActivationFunction=self.fusedActivationFunction,
padding=tflite.Padding.VALID,
)
else:
conv_op = tfl_ops.Conv3dOperator(
self.inputs,
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],
fusedActivationFunction=self.fusedActivationFunction,
padding=tflite.Padding.VALID,
)
ops = self.wrap_ops_with_nhwc_nchw_transposes([conv_op])
conv_op = ops[1]
# Pad handling
if sum(self.padding) > 0:
if weight_dim in (3, 4):
pad_h = self.padding[0]
pad_w = self.padding[1]
pad = [[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]
pad = [[0, 0], [pad_d, pad_d], [pad_h, pad_h], [pad_w, pad_w], [0, 0]]
pad_tensor = self.create_attr_tensor(np.array(pad, dtype='int32'))
pad_input = ops[0].outputs[0]
pad_array = np.pad(pad_input.tensor, pad)
pad_out = self.create_transform_tensor(pad_array, quantization=pad_input.quantization)
ops[1].inputs[0] = pad_out
pad_op = tfl_ops.PadOperator([pad_input, pad_tensor], [pad_out])
ops.insert(1, pad_op)
# Weight handling
weight = conv_op.inputs[1]
if conv_op.op.code == tflite.BuiltinOperator.DEPTHWISE_CONV_2D:
nchw2chwn_perm = np.array([1, 2, 3, 0], dtype='int32')
nchw2chwn_perm_tensor = self.create_attr_tensor(nchw2chwn_perm)
weight_q = weight.quantization
if weight_q is not None and weight_q.dim is not None:
new_dim = np.nonzero(nchw2chwn_perm == weight_q.dim)[0][0]
weight_q = QuantizationParameters(weight_q.scale, weight_q.zero_point, new_dim)
reordered_weight = self.create_transform_tensor(
np.transpose(weight.tensor, nchw2chwn_perm), quantization=weight_q
)
conv_op.inputs[1] = reordered_weight
reorder_op = tfl_ops.TransposeOperator([weight, nchw2chwn_perm_tensor], [reordered_weight])
else:
if weight_dim in (3, 4):
nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32')
nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)
else:
nchw2nhwc_perm = np.array([2, 3, 4, 1, 0], dtype='int32')
nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)
weight_q = weight.quantization
if weight_q is not None and weight_q.dim is not None:
new_dim = np.nonzero(nchw2nhwc_perm == weight_q.dim)[0][0]
weight_q = QuantizationParameters(weight_q.scale, weight_q.zero_point, new_dim)
reordered_weight = self.create_transform_tensor(
np.transpose(weight.tensor, nchw2nhwc_perm), quantization=weight_q
)
conv_op.inputs[1] = reordered_weight
reorder_op = tfl_ops.TransposeOperator([weight, nchw2nhwc_perm_tensor], [reordered_weight])
ops.insert(1, reorder_op)
# Bias handling
kernel_num = self.inputs[1].shape[0]
if conv_op.op.code in (tflite.BuiltinOperator.DEPTHWISE_CONV_2D, tflite.BuiltinOperator.CONV_3D):
kernel_num = self.inputs[1].shape[-1]
if len(conv_op.inputs) == 2 or conv_op.inputs[2] is None:
if conv_op.inputs[0].dtype == np.dtype('float32'):
bias = np.zeros((kernel_num,), dtype='float32')
q_args = None
else:
bias = np.zeros((kernel_num,), dtype='int32')
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))
elif conv_op.inputs[2].shape[0] != kernel_num and conv_op.inputs[2].shape[0] == 1:
if conv_op.inputs[0].dtype == np.float32:
bias = torch.tensor([conv_op.inputs[2][0]] * kernel_num, dtype='float32')
else:
bias = torch.tensor([conv_op.inputs[2][0]] * kernel_num, dtype='int32')
conv_op.inputs[2] = self.create_attr_tensor(bias)
ops = prev_ops + ops + next_ops
for op in ops:
graph_converter.add_operator(op, transform=True)
graph_converter.try_restore_edges(mapping)
for op in ops[:-1]:
output_name = op.outputs[0].name
node_name = graph_converter.tensor_node_map[output_name]
node = graph_converter.graph.vs.find(name=node_name)
assert node.outdegree() > 0, (
'The following node should be a part of the transformable node, but the outdegree of'
f' it is zero. {node}'
)
next_node = graph_converter.graph.vs[node.out_edges()[0].target]
assert next_node['node_type'] != ExtendedOperator.CONSTANT_NODE