blob: c25f4156c90337212828d5218a9e2bca9cbfa144 [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# TensorFlow Lite mapping functions used by both reader and writer.
18# Contains a mapping from the various TensorFlow Lite enums and options structs, generated by the FlatBuffer code
19# generator, to Vela's internal format.
Tim Hall79d07d22020-04-27 18:20:16 +010020import struct
21
Diego Russoea6111a2020-04-14 18:41:58 +010022import numpy as np
23
Tim Hall79d07d22020-04-27 18:20:16 +010024from .data_type import DataType
Diego Russoe8a10452020-04-21 17:39:10 +010025from .tflite import AbsOptions
26from .tflite import AddNOptions
Tim Hall79d07d22020-04-27 18:20:16 +010027from .tflite import AddOptions
Tim Hall79d07d22020-04-27 18:20:16 +010028from .tflite import ArgMaxOptions
Tim Hall79d07d22020-04-27 18:20:16 +010029from .tflite import ArgMinOptions
Jacob Bohlin8daf6b72020-09-15 16:28:35 +020030from .tflite import BatchMatMulOptions
Diego Russoe8a10452020-04-21 17:39:10 +010031from .tflite import BatchToSpaceNDOptions
Tim Hall79d07d22020-04-27 18:20:16 +010032from .tflite import BidirectionalSequenceLSTMOptions
33from .tflite import BidirectionalSequenceRNNOptions
Diego Russoe8a10452020-04-21 17:39:10 +010034from .tflite import CallOptions
35from .tflite import CastOptions
36from .tflite import ConcatEmbeddingsOptions
37from .tflite import ConcatenationOptions
38from .tflite import Conv2DOptions
Tim Hall79d07d22020-04-27 18:20:16 +010039from .tflite import CosOptions
Tim Hall79d07d22020-04-27 18:20:16 +010040from .tflite import DensifyOptions
41from .tflite import DepthToSpaceOptions
Diego Russoe8a10452020-04-21 17:39:10 +010042from .tflite import DepthwiseConv2DOptions
43from .tflite import DequantizeOptions
44from .tflite import DivOptions
45from .tflite import EmbeddingLookupSparseOptions
46from .tflite import EqualOptions
47from .tflite import ExpandDimsOptions
48from .tflite import ExpOptions
49from .tflite import FakeQuantOptions
50from .tflite import FillOptions
51from .tflite import FloorDivOptions
52from .tflite import FloorModOptions
53from .tflite import FullyConnectedOptions
54from .tflite import GatherNdOptions
55from .tflite import GatherOptions
56from .tflite import GreaterEqualOptions
57from .tflite import GreaterOptions
Louis Verhaard678645b2020-06-15 15:22:47 +020058from .tflite import HardSwishOptions
Tim Hall79d07d22020-04-27 18:20:16 +010059from .tflite import IfOptions
Diego Russoe8a10452020-04-21 17:39:10 +010060from .tflite import L2NormOptions
61from .tflite import LeakyReluOptions
62from .tflite import LessEqualOptions
63from .tflite import LessOptions
64from .tflite import LocalResponseNormalizationOptions
65from .tflite import LogicalAndOptions
66from .tflite import LogicalNotOptions
67from .tflite import LogicalOrOptions
68from .tflite import LogSoftmaxOptions
69from .tflite import LSHProjectionOptions
70from .tflite import LSTMOptions
71from .tflite import MatrixDiagOptions
72from .tflite import MatrixSetDiagOptions
73from .tflite import MaximumMinimumOptions
74from .tflite import MirrorPadOptions
75from .tflite import MulOptions
76from .tflite import NegOptions
Tim Hall79d07d22020-04-27 18:20:16 +010077from .tflite import NonMaxSuppressionV4Options
78from .tflite import NonMaxSuppressionV5Options
Diego Russoe8a10452020-04-21 17:39:10 +010079from .tflite import NotEqualOptions
80from .tflite import OneHotOptions
81from .tflite import PackOptions
82from .tflite import PadOptions
83from .tflite import PadV2Options
84from .tflite import Pool2DOptions
85from .tflite import PowOptions
86from .tflite import QuantizeOptions
87from .tflite import RangeOptions
88from .tflite import RankOptions
89from .tflite import ReducerOptions
90from .tflite import ReshapeOptions
91from .tflite import ResizeBilinearOptions
92from .tflite import ResizeNearestNeighborOptions
93from .tflite import ReverseSequenceOptions
94from .tflite import ReverseV2Options
95from .tflite import RNNOptions
Tim Hall79d07d22020-04-27 18:20:16 +010096from .tflite import ScatterNdOptions
97from .tflite import SegmentSumOptions
Diego Russoe8a10452020-04-21 17:39:10 +010098from .tflite import SelectOptions
Tim Hall79d07d22020-04-27 18:20:16 +010099from .tflite import SelectV2Options
Diego Russoe8a10452020-04-21 17:39:10 +0100100from .tflite import SequenceRNNOptions
101from .tflite import ShapeOptions
102from .tflite import SkipGramOptions
103from .tflite import SliceOptions
104from .tflite import SoftmaxOptions
105from .tflite import SpaceToBatchNDOptions
106from .tflite import SpaceToDepthOptions
107from .tflite import SparseToDenseOptions
108from .tflite import SplitOptions
109from .tflite import SplitVOptions
110from .tflite import SquaredDifferenceOptions
111from .tflite import SquareOptions
112from .tflite import SqueezeOptions
113from .tflite import StridedSliceOptions
114from .tflite import SubOptions
115from .tflite import SVDFOptions
116from .tflite import TileOptions
117from .tflite import TopKV2Options
118from .tflite import TransposeConvOptions
119from .tflite import TransposeOptions
120from .tflite import UnidirectionalSequenceLSTMOptions
121from .tflite import UniqueOptions
122from .tflite import UnpackOptions
123from .tflite import WhereOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100124from .tflite import WhileOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100125from .tflite import ZerosLikeOptions
126from .tflite.ActivationFunctionType import ActivationFunctionType
Diego Russoea6111a2020-04-14 18:41:58 +0100127from .tflite.BuiltinOperator import BuiltinOperator
128from .tflite.BuiltinOptions import BuiltinOptions
129from .tflite.Padding import Padding
Diego Russoe8a10452020-04-21 17:39:10 +0100130from .tflite.TensorType import TensorType
Tim Hall79d07d22020-04-27 18:20:16 +0100131
132
133def inverse_map(map):
134 return {v: k for k, v in map.items()}
135
136
137datatype_map = {
138 TensorType.UINT8: DataType.uint8,
139 TensorType.INT8: DataType.int8,
140 TensorType.INT16: DataType.int16,
141 TensorType.INT32: DataType.int32,
142 TensorType.INT64: DataType.int64,
143 TensorType.FLOAT16: DataType.float16,
144 TensorType.FLOAT32: DataType.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200145 TensorType.FLOAT64: DataType.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100146 TensorType.STRING: DataType.string,
147 TensorType.BOOL: DataType.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200148 TensorType.COMPLEX64: DataType.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200149 TensorType.COMPLEX128: DataType.complex128,
Tim Hall79d07d22020-04-27 18:20:16 +0100150}
151
152datatype_inv_map = inverse_map(datatype_map)
153datatype_inv_map[DataType.quint8] = TensorType.UINT8
154
155datatype_inv_map[DataType.qint8] = TensorType.INT8
156datatype_inv_map[DataType.qint16] = TensorType.INT16
157datatype_inv_map[DataType.qint32] = TensorType.INT32
158
159
160datatype_map_numpy = {
161 TensorType.UINT8: np.uint8,
162 TensorType.INT8: np.int8,
163 TensorType.INT16: np.int16,
164 TensorType.INT32: np.int32,
165 TensorType.INT64: np.int64,
166 TensorType.FLOAT16: np.float16,
167 TensorType.FLOAT32: np.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200168 TensorType.FLOAT64: np.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100169 TensorType.BOOL: np.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200170 TensorType.COMPLEX64: np.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200171 TensorType.COMPLEX128: np.complex128,
Tim Hall79d07d22020-04-27 18:20:16 +0100172}
173
174
175builtin_options_map = {
176 BuiltinOptions.Conv2DOptions: Conv2DOptions.Conv2DOptions,
177 BuiltinOptions.DepthwiseConv2DOptions: DepthwiseConv2DOptions.DepthwiseConv2DOptions,
178 BuiltinOptions.ConcatEmbeddingsOptions: ConcatEmbeddingsOptions.ConcatEmbeddingsOptions,
179 BuiltinOptions.LSHProjectionOptions: LSHProjectionOptions.LSHProjectionOptions,
180 BuiltinOptions.Pool2DOptions: Pool2DOptions.Pool2DOptions,
181 BuiltinOptions.SVDFOptions: SVDFOptions.SVDFOptions,
182 BuiltinOptions.RNNOptions: RNNOptions.RNNOptions,
183 BuiltinOptions.FullyConnectedOptions: FullyConnectedOptions.FullyConnectedOptions,
184 BuiltinOptions.SoftmaxOptions: SoftmaxOptions.SoftmaxOptions,
185 BuiltinOptions.ConcatenationOptions: ConcatenationOptions.ConcatenationOptions,
186 BuiltinOptions.AddOptions: AddOptions.AddOptions,
187 BuiltinOptions.L2NormOptions: L2NormOptions.L2NormOptions,
188 BuiltinOptions.LocalResponseNormalizationOptions: LocalResponseNormalizationOptions.LocalResponseNormalizationOptions, # noqa: E501
189 BuiltinOptions.LSTMOptions: LSTMOptions.LSTMOptions,
190 BuiltinOptions.ResizeBilinearOptions: ResizeBilinearOptions.ResizeBilinearOptions,
191 BuiltinOptions.CallOptions: CallOptions.CallOptions,
192 BuiltinOptions.ReshapeOptions: ReshapeOptions.ReshapeOptions,
193 BuiltinOptions.SkipGramOptions: SkipGramOptions.SkipGramOptions,
194 BuiltinOptions.SpaceToDepthOptions: SpaceToDepthOptions.SpaceToDepthOptions,
195 BuiltinOptions.EmbeddingLookupSparseOptions: EmbeddingLookupSparseOptions.EmbeddingLookupSparseOptions,
196 BuiltinOptions.MulOptions: MulOptions.MulOptions,
197 BuiltinOptions.PadOptions: PadOptions.PadOptions,
198 BuiltinOptions.GatherOptions: GatherOptions.GatherOptions,
199 BuiltinOptions.BatchToSpaceNDOptions: BatchToSpaceNDOptions.BatchToSpaceNDOptions,
200 BuiltinOptions.SpaceToBatchNDOptions: SpaceToBatchNDOptions.SpaceToBatchNDOptions,
201 BuiltinOptions.TransposeOptions: TransposeOptions.TransposeOptions,
202 BuiltinOptions.ReducerOptions: ReducerOptions.ReducerOptions,
203 BuiltinOptions.SubOptions: SubOptions.SubOptions,
204 BuiltinOptions.DivOptions: DivOptions.DivOptions,
205 BuiltinOptions.SqueezeOptions: SqueezeOptions.SqueezeOptions,
206 BuiltinOptions.SequenceRNNOptions: SequenceRNNOptions.SequenceRNNOptions,
207 BuiltinOptions.StridedSliceOptions: StridedSliceOptions.StridedSliceOptions,
208 BuiltinOptions.ExpOptions: ExpOptions.ExpOptions,
209 BuiltinOptions.TopKV2Options: TopKV2Options.TopKV2Options,
210 BuiltinOptions.SplitOptions: SplitOptions.SplitOptions,
211 BuiltinOptions.LogSoftmaxOptions: LogSoftmaxOptions.LogSoftmaxOptions,
212 BuiltinOptions.CastOptions: CastOptions.CastOptions,
213 BuiltinOptions.DequantizeOptions: DequantizeOptions.DequantizeOptions,
214 BuiltinOptions.MaximumMinimumOptions: MaximumMinimumOptions.MaximumMinimumOptions,
215 BuiltinOptions.ArgMaxOptions: ArgMaxOptions.ArgMaxOptions,
216 BuiltinOptions.LessOptions: LessOptions.LessOptions,
217 BuiltinOptions.NegOptions: NegOptions.NegOptions,
218 BuiltinOptions.PadV2Options: PadV2Options.PadV2Options,
219 BuiltinOptions.GreaterOptions: GreaterOptions.GreaterOptions,
220 BuiltinOptions.GreaterEqualOptions: GreaterEqualOptions.GreaterEqualOptions,
221 BuiltinOptions.LessEqualOptions: LessEqualOptions.LessEqualOptions,
222 BuiltinOptions.SelectOptions: SelectOptions.SelectOptions,
223 BuiltinOptions.SliceOptions: SliceOptions.SliceOptions,
224 BuiltinOptions.TransposeConvOptions: TransposeConvOptions.TransposeConvOptions,
225 BuiltinOptions.SparseToDenseOptions: SparseToDenseOptions.SparseToDenseOptions,
226 BuiltinOptions.TileOptions: TileOptions.TileOptions,
227 BuiltinOptions.ExpandDimsOptions: ExpandDimsOptions.ExpandDimsOptions,
228 BuiltinOptions.EqualOptions: EqualOptions.EqualOptions,
229 BuiltinOptions.NotEqualOptions: NotEqualOptions.NotEqualOptions,
230 BuiltinOptions.ShapeOptions: ShapeOptions.ShapeOptions,
231 BuiltinOptions.PowOptions: PowOptions.PowOptions,
232 BuiltinOptions.ArgMinOptions: ArgMinOptions.ArgMinOptions,
233 BuiltinOptions.FakeQuantOptions: FakeQuantOptions.FakeQuantOptions,
234 BuiltinOptions.PackOptions: PackOptions.PackOptions,
235 BuiltinOptions.LogicalOrOptions: LogicalOrOptions.LogicalOrOptions,
236 BuiltinOptions.OneHotOptions: OneHotOptions.OneHotOptions,
237 BuiltinOptions.LogicalAndOptions: LogicalAndOptions.LogicalAndOptions,
238 BuiltinOptions.LogicalNotOptions: LogicalNotOptions.LogicalNotOptions,
239 BuiltinOptions.UnpackOptions: UnpackOptions.UnpackOptions,
240 BuiltinOptions.FloorDivOptions: FloorDivOptions.FloorDivOptions,
241 BuiltinOptions.SquareOptions: SquareOptions.SquareOptions,
242 BuiltinOptions.ZerosLikeOptions: ZerosLikeOptions.ZerosLikeOptions,
243 BuiltinOptions.FillOptions: FillOptions.FillOptions,
244 BuiltinOptions.BidirectionalSequenceLSTMOptions: BidirectionalSequenceLSTMOptions.BidirectionalSequenceLSTMOptions,
245 BuiltinOptions.BidirectionalSequenceRNNOptions: BidirectionalSequenceRNNOptions.BidirectionalSequenceRNNOptions,
246 BuiltinOptions.UnidirectionalSequenceLSTMOptions: UnidirectionalSequenceLSTMOptions.UnidirectionalSequenceLSTMOptions, # noqa: E501
247 BuiltinOptions.FloorModOptions: FloorModOptions.FloorModOptions,
248 BuiltinOptions.RangeOptions: RangeOptions.RangeOptions,
249 BuiltinOptions.ResizeNearestNeighborOptions: ResizeNearestNeighborOptions.ResizeNearestNeighborOptions,
250 BuiltinOptions.LeakyReluOptions: LeakyReluOptions.LeakyReluOptions,
251 BuiltinOptions.SquaredDifferenceOptions: SquaredDifferenceOptions.SquaredDifferenceOptions,
252 BuiltinOptions.MirrorPadOptions: MirrorPadOptions.MirrorPadOptions,
253 BuiltinOptions.AbsOptions: AbsOptions.AbsOptions,
254 BuiltinOptions.SplitVOptions: SplitVOptions.SplitVOptions,
255 BuiltinOptions.UniqueOptions: UniqueOptions.UniqueOptions,
256 BuiltinOptions.ReverseV2Options: ReverseV2Options.ReverseV2Options,
257 BuiltinOptions.AddNOptions: AddNOptions.AddNOptions,
258 BuiltinOptions.GatherNdOptions: GatherNdOptions.GatherNdOptions,
259 BuiltinOptions.CosOptions: CosOptions.CosOptions,
260 BuiltinOptions.WhereOptions: WhereOptions.WhereOptions,
261 BuiltinOptions.RankOptions: RankOptions.RankOptions,
262 BuiltinOptions.ReverseSequenceOptions: ReverseSequenceOptions.ReverseSequenceOptions,
263 BuiltinOptions.MatrixDiagOptions: MatrixDiagOptions.MatrixDiagOptions,
264 BuiltinOptions.QuantizeOptions: QuantizeOptions.QuantizeOptions,
265 BuiltinOptions.MatrixSetDiagOptions: MatrixSetDiagOptions.MatrixSetDiagOptions,
266 BuiltinOptions.DensifyOptions: DensifyOptions.DensifyOptions,
267 BuiltinOptions.DepthToSpaceOptions: DepthToSpaceOptions.DepthToSpaceOptions,
Louis Verhaard678645b2020-06-15 15:22:47 +0200268 BuiltinOptions.HardSwishOptions: HardSwishOptions.HardSwishOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100269 BuiltinOptions.IfOptions: IfOptions.IfOptions,
270 BuiltinOptions.NonMaxSuppressionV4Options: NonMaxSuppressionV4Options.NonMaxSuppressionV4Options,
271 BuiltinOptions.NonMaxSuppressionV5Options: NonMaxSuppressionV5Options.NonMaxSuppressionV5Options,
272 BuiltinOptions.ScatterNdOptions: ScatterNdOptions.ScatterNdOptions,
273 BuiltinOptions.SegmentSumOptions: SegmentSumOptions.SegmentSumOptions,
274 BuiltinOptions.SelectV2Options: SelectV2Options.SelectV2Options,
275 BuiltinOptions.WhileOptions: WhileOptions.WhileOptions,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200276 BuiltinOptions.BatchMatMulOptions: BatchMatMulOptions.BatchMatMulOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100277}
278
279builtin_options_inv_map = inverse_map(builtin_options_map)
280
281
282def underscore_to_camel_case(s):
283 return "".join(x.title() for x in s.split("_"))
284
285
286def padding_deserialize(x):
287 return padding_map[x]
288
289
290def padding_serialize(builder, x):
291 return padding_inv_map[x]
292
293
294def activation_deserialize(x):
295 return activation_function_map[x]
296
297
298def activation_serialize(builder, x):
299 return activation_function_inv_map[x]
300
301
302def datatype_deserialize(x):
303 return datatype_map[x]
304
305
306def datatype_serialize(builder, x):
307 return datatype_inv_map[x]
308
309
310def identity(x):
311 return x
312
313
314def identity_serialize(builder, x):
315 return x
316
317
318def write_byte_vector(builder, v):
319 builder.StartVector(1, len(v), 1)
320 for e in v[::-1]:
321 builder.PrependByte(e)
322 return builder.EndVector(len(v))
323
324
325def write_int_vector(builder, v):
326 builder.StartVector(4, len(v), 4)
327 for e in v[::-1]:
328 builder.PrependInt32(e)
329 return builder.EndVector(len(v))
330
331
332class OptionsSerializer:
333 def __init__(self, name, members=[]):
334 self.name = name
335 self.module = globals()[self.name]
336 self.cls = getattr(self.module, self.name)
337 self.builtin_opt_type = builtin_options_inv_map[self.cls]
Tim Hall79d07d22020-04-27 18:20:16 +0100338 self.members = []
339 for mem in members:
340 deserialize = identity
341 serialize = identity_serialize
342 is_vector = False
343 if isinstance(mem, tuple):
344 if len(mem) == 3:
345 mem, deserialize, serialize = mem
346 elif len(mem) == 2:
347 mem, is_vector = mem
348 deserialize = tuple
349 serialize = write_int_vector
350 else:
351 assert 0
352 underscore_mem = mem
353 camelcase_mem = underscore_to_camel_case(mem)
354 self.members.append((underscore_mem, camelcase_mem, deserialize, serialize, is_vector))
355
Tim Hallc8310b12020-06-17 14:53:11 +0100356 def deserialize(self, op_data):
357 builtin_options = op_data.BuiltinOptions()
Tim Hall79d07d22020-04-27 18:20:16 +0100358 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100359 if builtin_options:
Tim Hall79d07d22020-04-27 18:20:16 +0100360 tfattrs = self.cls()
Tim Hallc8310b12020-06-17 14:53:11 +0100361 tfattrs.Init(builtin_options.Bytes, builtin_options.Pos)
Tim Hall79d07d22020-04-27 18:20:16 +0100362 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
363 fun = camelcase_mem
364 if is_vector:
365 fun += "AsNumpy"
366
367 a = deserialize(getattr(tfattrs, fun)())
368 attrs[underscore_mem] = a
369 return attrs
370
371 def serialize(self, builder, attrs):
372 ser_attrs = []
373 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
374 a = serialize(builder, attrs[underscore_mem])
375 ser_attrs.append((camelcase_mem, a))
376
377 getattr(self.module, self.name + "Start")(builder)
378
379 for camelcase_mem, a in ser_attrs:
380 getattr(self.module, self.name + "Add" + camelcase_mem)(builder, a)
381
382 return getattr(self.module, self.name + "End")(builder), None
383
384
385class CustomOptionsSerializer:
Tim Hallc8310b12020-06-17 14:53:11 +0100386 CUSTOM_OPTIONS_NPU_OP = [0x01, 0x04, 0x01] # NpuOp=1, FlexbufferFormat.UINT8=4, byte length=1
387 CUSTOM_OPTIONS_FORMAT_DEFAULT = 0
388
Tim Hall79d07d22020-04-27 18:20:16 +0100389 def __init__(self):
Tim Hall79d07d22020-04-27 18:20:16 +0100390 self.custom_opt_format = 0
391
Tim Hallc8310b12020-06-17 14:53:11 +0100392 def deserialize(self, op_data):
Tim Hall79d07d22020-04-27 18:20:16 +0100393 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100394 custom_options = op_data.CustomOptionsAsNumpy()
395 attrs["custom_options"] = custom_options
396 attrs["custom_options_format"] = op_data.CustomOptionsFormat()
397
398 if np.array_equal(custom_options, self.CUSTOM_OPTIONS_NPU_OP):
399 attrs["custom_type"] = "ExistingNpuOp"
400
Tim Hall79d07d22020-04-27 18:20:16 +0100401 return attrs
402
403 def serialize(self, builder, attrs):
Tim Hallc8310b12020-06-17 14:53:11 +0100404 custom_type = attrs.get("custom_type", "")
405 self.custom_opt_format = attrs.get("custom_options_format", self.CUSTOM_OPTIONS_FORMAT_DEFAULT)
Tim Hall79d07d22020-04-27 18:20:16 +0100406
407 # Set NPU op custom options for the TensorFlow Lite custom operator
Tim Hallc8310b12020-06-17 14:53:11 +0100408 if custom_type == "NpuOp":
409 custom_options = self.CUSTOM_OPTIONS_NPU_OP
410 else:
411 custom_options = attrs.get("custom_options", [])
Tim Hall79d07d22020-04-27 18:20:16 +0100412
Tim Hallc8310b12020-06-17 14:53:11 +0100413 custom_options_bytes = struct.pack("<{0}B".format(len(custom_options)), *custom_options)
414 custom_offset = write_byte_vector(builder, custom_options_bytes)
Tim Hall79d07d22020-04-27 18:20:16 +0100415
416 return None, custom_offset
417
418
419padding_map = {
420 Padding.SAME: b"SAME",
421 Padding.VALID: b"VALID",
422}
423
424padding_inv_map = inverse_map(padding_map)
425
426
427activation_function_map = {
428 ActivationFunctionType.NONE: None,
429 ActivationFunctionType.RELU: "Relu",
430 ActivationFunctionType.RELU_N1_TO_1: "ReluN1To1",
431 ActivationFunctionType.RELU6: "Relu6",
432 ActivationFunctionType.TANH: "Tanh",
433 ActivationFunctionType.SIGN_BIT: "SignBit",
434}
435
436activation_function_inv_map = inverse_map(activation_function_map)
437
438fused_act = ("fused_activation_function", activation_deserialize, activation_serialize)
439padding = ("padding", padding_deserialize, padding_serialize)
440
441pool2d_opts = OptionsSerializer(
442 "Pool2DOptions", (padding, "stride_w", "stride_h", "filter_width", "filter_height", fused_act,)
443)
444
445depthwise_opts = OptionsSerializer(
446 "DepthwiseConv2DOptions",
447 (padding, "stride_w", "stride_h", "depth_multiplier", fused_act, "dilation_w_factor", "dilation_h_factor",),
448)
449
450conv2d_opts = OptionsSerializer(
451 "Conv2DOptions", (padding, "stride_w", "stride_h", fused_act, "dilation_w_factor", "dilation_h_factor",)
452)
453
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200454lstm_opts = OptionsSerializer(
455 "LSTMOptions", (fused_act, "cell_clip", "proj_clip", "kernel_type", "asymmetric_quantize_inputs")
456)
Tim Hall79d07d22020-04-27 18:20:16 +0100457
458unidir_seq_lstm_opts = OptionsSerializer(
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200459 "UnidirectionalSequenceLSTMOptions",
460 (fused_act, "cell_clip", "proj_clip", "time_major", "asymmetric_quantize_inputs",),
Tim Hall79d07d22020-04-27 18:20:16 +0100461)
462
463bidir_seq_lstm_opts = OptionsSerializer(
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200464 "BidirectionalSequenceLSTMOptions",
465 (fused_act, "cell_clip", "proj_clip", "merge_outputs", "time_major", "asymmetric_quantize_inputs"),
Tim Hall79d07d22020-04-27 18:20:16 +0100466)
467
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200468rnn_opts = OptionsSerializer("RNNOptions", (fused_act, "asymmetric_quantize_inputs"))
Tim Hall79d07d22020-04-27 18:20:16 +0100469
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200470seq_rnn_opts = OptionsSerializer("SequenceRNNOptions", ("time_major", fused_act, "asymmetric_quantize_inputs",))
Tim Hall79d07d22020-04-27 18:20:16 +0100471
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200472bidir_seq_rnn_opts = OptionsSerializer(
473 "BidirectionalSequenceRNNOptions", ("time_major", fused_act, "merge_outputs", "asymmetric_quantize_inputs")
474)
Tim Hall79d07d22020-04-27 18:20:16 +0100475
476
477reducer_opts = OptionsSerializer("ReducerOptions", ("keep_dims",))
478
479is_int_vec = True
480
481custom_prefix = "Custom_"
482
483builtin_operator_map = {
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200484 BuiltinOperator.ADD: ("AddAct", OptionsSerializer("AddOptions", (fused_act, "pot_scale_int16"))),
Tim Hall79d07d22020-04-27 18:20:16 +0100485 BuiltinOperator.AVERAGE_POOL_2D: ("AvgPoolAct", pool2d_opts),
486 BuiltinOperator.CONCATENATION: ("ConcatTFLite", OptionsSerializer("ConcatenationOptions", ("axis", fused_act))),
487 BuiltinOperator.CONV_2D: ("Conv2DBiasAct", conv2d_opts),
488 BuiltinOperator.DEPTHWISE_CONV_2D: ("DepthwiseConv2dBiasAct", depthwise_opts),
489 BuiltinOperator.DEPTH_TO_SPACE: ("DepthToSpace", OptionsSerializer("DepthToSpaceOptions", ("block_size",))),
490 BuiltinOperator.DEQUANTIZE: ("Dequantize", OptionsSerializer("DequantizeOptions")),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200491 BuiltinOperator.EMBEDDING_LOOKUP: ("EmbeddingLookup", None),
Tim Hall79d07d22020-04-27 18:20:16 +0100492 BuiltinOperator.FLOOR: ("Floor", None),
493 BuiltinOperator.FULLY_CONNECTED: (
494 "FullyConnectedAct",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200495 OptionsSerializer("FullyConnectedOptions", (fused_act, "weights_format", "asymmetric_quantize_inputs")),
Tim Hall79d07d22020-04-27 18:20:16 +0100496 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200497 BuiltinOperator.HASHTABLE_LOOKUP: ("HashtableLookup", None),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200498 BuiltinOperator.L2_NORMALIZATION: ("L2NormAct", OptionsSerializer("L2NormOptions", (fused_act,))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200499 BuiltinOperator.L2_POOL_2D: ("L2Pool2D", pool2d_opts),
Tim Hall79d07d22020-04-27 18:20:16 +0100500 BuiltinOperator.LOCAL_RESPONSE_NORMALIZATION: (
501 "LRN",
502 OptionsSerializer("LocalResponseNormalizationOptions", ("radius", "bias", "alpha", "beta")),
503 ),
504 BuiltinOperator.LOGISTIC: ("Sigmoid", None),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200505 BuiltinOperator.LSH_PROJECTION: ("LSHProjection", OptionsSerializer("LSHProjectionOptions", ("type",))),
Tim Hall79d07d22020-04-27 18:20:16 +0100506 BuiltinOperator.LSTM: ("LstmAct", lstm_opts),
507 BuiltinOperator.MAX_POOL_2D: ("MaxPool", pool2d_opts),
508 BuiltinOperator.MUL: ("MulAct", OptionsSerializer("MulOptions", (fused_act,))),
509 BuiltinOperator.RELU: ("Relu", None),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200510 BuiltinOperator.RELU_N1_TO_1: ("ReluN1To1", None),
Tim Hall79d07d22020-04-27 18:20:16 +0100511 BuiltinOperator.RELU6: ("Relu6", None),
512 BuiltinOperator.RESHAPE: ("Reshape", OptionsSerializer("ReshapeOptions", (("new_shape", is_int_vec),))),
513 BuiltinOperator.RESIZE_BILINEAR: (
514 "ResizeBilinear",
515 OptionsSerializer("ResizeBilinearOptions", ("align_corners", "half_pixel_centers")),
516 ),
517 BuiltinOperator.RNN: ("RnnAct", rnn_opts),
518 BuiltinOperator.SOFTMAX: ("Softmax", OptionsSerializer("SoftmaxOptions", ("beta",))),
519 BuiltinOperator.SPACE_TO_DEPTH: ("SpaceToDepth", OptionsSerializer("SpaceToDepthOptions", ("block_size",))),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200520 BuiltinOperator.SVDF: (
521 "SvdfAct",
522 OptionsSerializer("SVDFOptions", ("rank", fused_act, "asymmetric_quantize_inputs")),
523 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100524 BuiltinOperator.TANH: ("Tanh", None),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200525 BuiltinOperator.CONCAT_EMBEDDINGS: (
526 "ConcatEmbeddings",
527 OptionsSerializer(
528 "ConcatEmbeddingsOptions",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200529 (
530 "num_channels",
531 "num_columns_per_channel",
532 "num_columns_per_channel_as_numpy",
533 "num_columns_per_channel_as_length",
534 "embedding_dim_per_channel",
535 "embedding_dim_per_channel_as_numpy",
536 "embedding_dim_per_channel_as_length",
537 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200538 ),
539 ),
540 BuiltinOperator.SKIP_GRAM: (
541 "SkipGram",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200542 OptionsSerializer("SkipGramOptions", ("ngram_size", "max_skip_size", "include_all_ngrams")),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200543 ),
544 BuiltinOperator.CALL: ("Call", OptionsSerializer("CallOptions", ("subgraph",))),
545 BuiltinOperator.EMBEDDING_LOOKUP_SPARSE: (
546 "EmbeddingLookupSparse",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200547 OptionsSerializer("EmbeddingLookupSparseOptions", ("combiner",)),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200548 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100549 BuiltinOperator.PAD: ("Pad", OptionsSerializer("PadOptions")),
550 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_RNN: ("UnidirectionalSequenceRnnAct", seq_rnn_opts),
551 BuiltinOperator.GATHER: ("GatherV2", OptionsSerializer("GatherOptions", ("axis",))),
552 BuiltinOperator.BATCH_TO_SPACE_ND: ("BatchToSpaceND", OptionsSerializer("BatchToSpaceNDOptions")),
553 BuiltinOperator.SPACE_TO_BATCH_ND: ("SpaceToBatchND", OptionsSerializer("SpaceToBatchNDOptions")),
554 BuiltinOperator.TRANSPOSE: ("Transpose", OptionsSerializer("TransposeOptions")),
555 BuiltinOperator.MEAN: ("Mean", None),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200556 BuiltinOperator.SUB: ("SubAct", OptionsSerializer("SubOptions", (fused_act, "pot_scale_int16",))),
Tim Hall79d07d22020-04-27 18:20:16 +0100557 BuiltinOperator.DIV: ("DivAct", OptionsSerializer("DivOptions", (fused_act,))),
558 BuiltinOperator.SQUEEZE: ("Squeeze", OptionsSerializer("SqueezeOptions", (("squeeze_dims", is_int_vec),))),
559 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_LSTM: ("UnidirectionalSequenceLstmAct", unidir_seq_lstm_opts),
560 BuiltinOperator.STRIDED_SLICE: (
561 "StridedSlice",
562 OptionsSerializer(
563 "StridedSliceOptions", ("begin_mask", "end_mask", "ellipsis_mask", "new_axis_mask", "shrink_axis_mask")
564 ),
565 ),
566 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_RNN: ("BidirectionalSequenceRnnAct", bidir_seq_rnn_opts),
567 BuiltinOperator.EXP: ("Exp", OptionsSerializer("ExpOptions")),
568 BuiltinOperator.TOPK_V2: ("TopKV2", OptionsSerializer("TopKV2Options")),
569 BuiltinOperator.SPLIT: ("Split", OptionsSerializer("SplitOptions", ("num_splits",))),
570 BuiltinOperator.LOG_SOFTMAX: ("LogSoftmax", OptionsSerializer("LogSoftmaxOptions")),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200571 BuiltinOperator.DELEGATE: ("Delegate", None),
Tim Hall79d07d22020-04-27 18:20:16 +0100572 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_LSTM: ("BidirectionalSequenceLstmAct", bidir_seq_lstm_opts),
573 BuiltinOperator.CAST: (
574 "Cast",
575 OptionsSerializer(
576 "CastOptions",
577 (
578 ("in_data_type", datatype_deserialize, datatype_serialize),
579 ("out_data_type", datatype_deserialize, datatype_serialize),
580 ),
581 ),
582 ),
Tim Hallc30f4952020-06-15 20:47:35 +0100583 BuiltinOperator.PRELU: ("Prelu", None),
Tim Hall79d07d22020-04-27 18:20:16 +0100584 BuiltinOperator.MAXIMUM: ("Maximum", OptionsSerializer("MaximumMinimumOptions")),
585 BuiltinOperator.ARG_MAX: (
586 "ArgMax",
587 OptionsSerializer("ArgMaxOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
588 ),
589 BuiltinOperator.MINIMUM: ("Minimum", OptionsSerializer("MaximumMinimumOptions")),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200590 BuiltinOperator.LESS: ("Less", OptionsSerializer("LessOptions")),
591 BuiltinOperator.NEG: ("Neg", OptionsSerializer("NegOptions")),
592 BuiltinOperator.PADV2: ("PadV2", OptionsSerializer("PadV2Options")),
593 BuiltinOperator.GREATER: ("Greater", OptionsSerializer("GreaterOptions")),
594 BuiltinOperator.GREATER_EQUAL: ("GreaterEqual", OptionsSerializer("GreaterEqualOptions")),
595 BuiltinOperator.LESS_EQUAL: ("LessEqual", OptionsSerializer("LessEqualOptions")),
596 BuiltinOperator.SELECT: ("Select", OptionsSerializer("SelectOptions")),
597 BuiltinOperator.SLICE: ("Slice", OptionsSerializer("SliceOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100598 BuiltinOperator.SIN: ("Sin", None),
599 BuiltinOperator.TRANSPOSE_CONV: (
600 "Conv2DBackpropInput",
601 OptionsSerializer("TransposeConvOptions", (padding, "stride_w", "stride_h")),
602 ),
603 BuiltinOperator.SPARSE_TO_DENSE: (
604 "SparseToDense",
605 OptionsSerializer("SparseToDenseOptions", ("validate_indices",)),
606 ),
607 BuiltinOperator.TILE: ("Tile", OptionsSerializer("TileOptions")),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200608 BuiltinOperator.EXPAND_DIMS: ("ExpandDims", OptionsSerializer("ExpandDimsOptions")),
609 BuiltinOperator.EQUAL: ("Equal", OptionsSerializer("EqualOptions")),
610 BuiltinOperator.NOT_EQUAL: ("NotEqual", OptionsSerializer("NotEqualOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100611 BuiltinOperator.LOG: ("Log", None),
612 BuiltinOperator.SUM: ("Sum", None),
613 BuiltinOperator.SQRT: ("Sqrt", None),
614 BuiltinOperator.RSQRT: ("Rsqrt", None),
615 BuiltinOperator.SHAPE: (
616 "Shape",
617 OptionsSerializer("ShapeOptions", (("out_type", datatype_deserialize, datatype_serialize),)),
618 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200619 BuiltinOperator.POW: ("Pow", OptionsSerializer("PowOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100620 BuiltinOperator.ARG_MIN: (
621 "ArgMin",
622 OptionsSerializer("ArgMinOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
623 ),
624 BuiltinOperator.FAKE_QUANT: (
625 "FakeQuantWithMinMaxArgs",
626 OptionsSerializer("FakeQuantOptions", ("min", "max", "num_bits", "narrow_range")),
627 ),
628 BuiltinOperator.REDUCE_PROD: ("Prod", reducer_opts),
629 BuiltinOperator.REDUCE_MAX: ("Max", reducer_opts),
630 BuiltinOperator.PACK: ("Pack", OptionsSerializer("PackOptions", ("values_count", "axis"))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200631 BuiltinOperator.LOGICAL_OR: ("LogicalOr", OptionsSerializer("LogicalOrOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100632 BuiltinOperator.ONE_HOT: ("OneHot", OptionsSerializer("OneHotOptions", ("axis",))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200633 BuiltinOperator.LOGICAL_AND: ("LogicalAnd", OptionsSerializer("LogicalAndOptions")),
634 BuiltinOperator.LOGICAL_NOT: ("LogicalNot", OptionsSerializer("LogicalNotOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100635 BuiltinOperator.UNPACK: ("Unpack", OptionsSerializer("UnpackOptions", ("num", "axis"))),
636 BuiltinOperator.REDUCE_MIN: ("Min", reducer_opts),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200637 BuiltinOperator.FLOOR_DIV: ("FloorDiv", OptionsSerializer("FloorDivOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100638 BuiltinOperator.REDUCE_ANY: ("Any", reducer_opts),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200639 BuiltinOperator.SQUARE: ("Square", OptionsSerializer("SquareOptions")),
640 BuiltinOperator.ZEROS_LIKE: ("ZerosLike", OptionsSerializer("ZerosLikeOptions")),
641 BuiltinOperator.FILL: ("Fill", OptionsSerializer("FillOptions")),
642 BuiltinOperator.FLOOR_MOD: ("FloorMod", OptionsSerializer("FloorModOptions")),
643 BuiltinOperator.RANGE: ("Range", OptionsSerializer("RangeOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100644 BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: (
645 "ResizeNearestNeighbor",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200646 OptionsSerializer("ResizeNearestNeighborOptions", ("align_corners", "half_pixel_centers")),
Tim Hall79d07d22020-04-27 18:20:16 +0100647 ),
648 BuiltinOperator.LEAKY_RELU: ("LeakyRelu", OptionsSerializer("LeakyReluOptions", ("alpha",))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200649 BuiltinOperator.SQUARED_DIFFERENCE: ("SquaredDifference", OptionsSerializer("SquaredDifferenceOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100650 BuiltinOperator.MIRROR_PAD: ("MirrorPad", OptionsSerializer("MirrorPadOptions", ("mode",))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200651 BuiltinOperator.ABS: ("Abs", OptionsSerializer("AbsOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100652 BuiltinOperator.SPLIT_V: ("SplitV", OptionsSerializer("SplitVOptions", ("num_splits",))),
653 BuiltinOperator.UNIQUE: (
654 "Unique",
655 OptionsSerializer("UniqueOptions", (("idx_out_type", datatype_deserialize, datatype_serialize),)),
656 ),
657 BuiltinOperator.CEIL: ("Ceil", None),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200658 BuiltinOperator.REVERSE_V2: ("ReverseV2", OptionsSerializer("ReverseV2Options")),
659 BuiltinOperator.ADD_N: ("AddN", OptionsSerializer("AddNOptions")),
660 BuiltinOperator.GATHER_ND: ("GatherNd", OptionsSerializer("GatherNdOptions")),
661 BuiltinOperator.COS: ("Cos", OptionsSerializer("CosOptions")),
662 BuiltinOperator.WHERE: ("Where", OptionsSerializer("WhereOptions")),
663 BuiltinOperator.RANK: ("Rank", OptionsSerializer("RankOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100664 BuiltinOperator.ELU: ("Elu", None),
665 BuiltinOperator.REVERSE_SEQUENCE: (
666 "ReverseSequence",
667 OptionsSerializer("ReverseSequenceOptions", ("seq_dim", "batch_dim")),
668 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200669 BuiltinOperator.MATRIX_DIAG: ("MatrixDiag", OptionsSerializer("MatrixDiagOptions")),
670 BuiltinOperator.QUANTIZE: ("Quantize", OptionsSerializer("QuantizeOptions")),
671 BuiltinOperator.MATRIX_SET_DIAG: ("MatrixSetDiag", OptionsSerializer("MatrixSetDiagOptions")),
Dwight Lidmana3aad932020-08-19 15:53:46 +0200672 BuiltinOperator.ROUND: ("Round", None),
Louis Verhaard678645b2020-06-15 15:22:47 +0200673 BuiltinOperator.HARD_SWISH: ("HardSwish", OptionsSerializer("HardSwishOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100674 BuiltinOperator.IF: ("If", OptionsSerializer("IfOptions", ("then_subgraph_index", "else_subgraph_index"))),
675 BuiltinOperator.WHILE: ("While", OptionsSerializer("WhileOptions", ("cond_subgraph_index", "body_subgraph_index"))),
676 BuiltinOperator.NON_MAX_SUPPRESSION_V4: ("NonMaxSuppressionV4", OptionsSerializer("NonMaxSuppressionV4Options")),
677 BuiltinOperator.NON_MAX_SUPPRESSION_V5: ("NonMaxSuppressionV5", OptionsSerializer("NonMaxSuppressionV5Options")),
678 BuiltinOperator.SCATTER_ND: ("ScatterNd", OptionsSerializer("ScatterNdOptions")),
679 BuiltinOperator.SELECT_V2: ("SelectV2", OptionsSerializer("SelectV2Options")),
680 BuiltinOperator.DENSIFY: ("Densify", OptionsSerializer("DensifyOptions")),
681 BuiltinOperator.SEGMENT_SUM: ("SegmentSum", OptionsSerializer("SegmentSumOptions")),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200682 BuiltinOperator.BATCH_MATMUL: ("BatchMatMul", OptionsSerializer("BatchMatMulOptions", ("adj_x", "adj_y"))),
Tim Hall79d07d22020-04-27 18:20:16 +0100683 BuiltinOperator.CUSTOM: (custom_prefix, CustomOptionsSerializer()),
684}
685
686builtin_operator_inv_map = {v[0]: (k, v[1]) for k, v in builtin_operator_map.items()}
687
688builtin_operator_inv_map["NpuOp"] = (BuiltinOperator.CUSTOM, CustomOptionsSerializer())