MLBEDSW-3148: Refactor Operation
- op.type is now an enum instead of a string
- Removed unused operator codes
- Refactored some attributes like npu_block_type, fused_activation_function
- Refactored operator index calculation
- Refactored a number of operator sets
Change-Id: I641f65ee375794b7aec42abc0664251ae37d78e8
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 68af487..f444ee5 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -22,6 +22,7 @@
from flatbuffers.builder import UOffsetTFlags
from .nn_graph import PassPlacement
+from .operation import Op
from .tensor import MemType
from .tensor import TensorPurpose
from .tflite import Buffer
@@ -34,7 +35,6 @@
from .tflite import Tensor
from .tflite_mapping import builtin_operator_inv_map
from .tflite_mapping import BuiltinOperator
-from .tflite_mapping import custom_prefix
from .tflite_mapping import datatype_inv_map
# ugh, the python flatbuffer interface is missing a method to add in file identifier. patching it in here:
@@ -77,7 +77,7 @@
self.scratch_fast_buf_id = 1 # Always assign scratch_fast to buffer 1
self.buffers_to_write = [] # have an empty array there
- self.ops_to_ignore = set(("Const", "Placeholder", "SubgraphInput"))
+ self.ops_to_ignore = set((Op.Const, Op.Placeholder, Op.SubgraphInput))
self.tensors_to_reshape = {}
@@ -89,16 +89,17 @@
for op in ps.ops:
if op.type not in self.ops_to_ignore:
all_ops.append(op)
- if op.type.startswith("Conv2D") or op.type.startswith("DepthwiseConv2d"):
+ if op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op():
# If values are None op has non-constant weights
if op.inputs[1].values is not None:
self.tensors_to_reshape[op.inputs[1]] = (3, 0, 1, 2)
- if op.type.startswith("FullyConnected"):
+ if op.type == Op.FullyConnected:
# If values are None op has non-constant weights
if op.inputs[1].values is not None:
self.tensors_to_reshape[op.inputs[1]] = (1, 0)
- self.operator_codes = list(sorted(set(op.type for op in all_ops)))
+ # list of tuple(Op, string); the custom code is only used for 3rd party custom operators
+ self.operator_codes = sorted(set((op.type, op.attrs.get("custom_code", "")) for op in all_ops))
self.operator_code_map = {}
def write_byte_vector(self, v, alignment=1):
@@ -163,25 +164,25 @@
return buffer_map
- def serialise_operator_code(self, idx, code):
+ def serialise_operator_code(self, idx, op_type, custom_code):
builder = self.builder
custom_code_offset = None
- if code.startswith(custom_prefix):
- tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
- custom_code_offset = builder.CreateString(code[len(custom_prefix) :])
+ if op_type == Op.Custom:
+ tf_code, opt_serializer = builtin_operator_inv_map[op_type]
+ custom_code_offset = builder.CreateString(custom_code)
else:
assert (
- code in builtin_operator_inv_map
- ), "Vela does not contain a mapping to serialise {} operator to a TensorFlow Lite operator".format(code)
- tf_code, opt_serializer = builtin_operator_inv_map[code]
+ op_type in builtin_operator_inv_map
+ ), "Vela does not contain a mapping to serialise {} operator to a TensorFlow Lite operator".format(op_type)
+ tf_code, opt_serializer = builtin_operator_inv_map[op_type]
if tf_code == BuiltinOperator.CUSTOM:
assert (
- code == "NpuOp"
+ op_type == Op.CustomNpuOp
), "Vela only supports serialising NpuOp operators as TensorFlow Lite Custom operators"
custom_code_offset = builder.CreateString("ethos-u")
- self.operator_code_map[code] = (idx, tf_code, opt_serializer)
+ self.operator_code_map[op_type] = (idx, tf_code, opt_serializer)
OperatorCode.OperatorCodeStart(builder)
OperatorCode.OperatorCodeAddBuiltinCode(builder, tf_code)
@@ -281,6 +282,8 @@
attrs["dilation_w_factor"] = attrs["dilation"][2]
if "channel_multiplier" in attrs:
attrs["depth_multiplier"] = attrs["channel_multiplier"]
+ if op.activation is not None:
+ attrs["fused_activation_function"] = op.activation
builtin_opt_offset, custom_opt_offset = opt_serializer.serialize(builder, attrs)
@@ -310,7 +313,7 @@
for op in ps.ops:
if op.type not in self.ops_to_ignore:
all_ops.append(op)
- elif op.type == "Placeholder":
+ elif op.type == Op.Placeholder:
placeholder_ops.append(op)
# Add the tensors from all valid ops, as well as the tensors from placeholder ops
@@ -404,7 +407,7 @@
def serialise_model(self):
builder = self.builder
operator_code_offset = self.write_offset_vector(
- [self.serialise_operator_code(idx, code) for idx, code in enumerate(self.operator_codes)]
+ [self.serialise_operator_code(idx, optype, code) for idx, (optype, code) in enumerate(self.operator_codes)]
)
description = builder.CreateString("Vela Optimised")