128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
""" Tensor class """
|
|
|
|
from typing import Dict, Optional
|
|
import string
|
|
import textwrap
|
|
|
|
from tflite_micro.codegen import utils
|
|
from tflite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
|
|
|
|
_TENSOR_TYPES: Dict[int, str] = {
|
|
schema_fb.TensorType.FLOAT16: "kTfLiteFloat16",
|
|
schema_fb.TensorType.FLOAT32: "kTfLiteFloat32",
|
|
schema_fb.TensorType.FLOAT64: "kTfLiteFloat64",
|
|
schema_fb.TensorType.INT16: "kTfLiteInt16",
|
|
schema_fb.TensorType.UINT16: "kTfLiteUInt16",
|
|
schema_fb.TensorType.INT32: "kTfLiteInt32",
|
|
schema_fb.TensorType.UINT32: "kTfLiteUInt32",
|
|
schema_fb.TensorType.UINT8: "kTfLiteUInt8",
|
|
schema_fb.TensorType.INT8: "kTfLiteInt8",
|
|
schema_fb.TensorType.INT64: "kTfLiteInt64",
|
|
schema_fb.TensorType.UINT64: "kTfLiteUInt64",
|
|
schema_fb.TensorType.STRING: "kTfLiteString",
|
|
schema_fb.TensorType.BOOL: "kTfLiteBool",
|
|
schema_fb.TensorType.COMPLEX64: "kTfLiteComplex64",
|
|
schema_fb.TensorType.COMPLEX128: "kTfLiteComplex128",
|
|
schema_fb.TensorType.RESOURCE: "kTfLiteResource",
|
|
schema_fb.TensorType.VARIANT: "kTfLiteVariant",
|
|
schema_fb.TensorType.INT4: "kTfLiteInt4",
|
|
}
|
|
|
|
|
|
class Buffer(object):
|
|
""" This buffer could be either a static array or a pointer into the arena """
|
|
|
|
def __init__(self, buffer_name: str, buffer: schema_fb.BufferT):
|
|
# TODO(rjascani): Get arena allocation offsets from preprocessor
|
|
self._buffer_name = buffer_name
|
|
self._buffer = buffer
|
|
|
|
@property
|
|
def address(self) -> str:
|
|
if self._buffer is None or self._buffer.data is None:
|
|
# TODO(rjascani): This needs to point into the arena
|
|
return f"nullptr /* {self._buffer_name} */"
|
|
return f"&{self._buffer_name}"
|
|
|
|
def generate_c_buffer_array(self, indent: str) -> str:
|
|
if self._buffer is None or self._buffer.data is None:
|
|
return f"// {self._buffer_name} is located in the arena\n"
|
|
|
|
buffer_template = string.Template(
|
|
"alignas(16) uint8_t ${buffer_name}[${size}] = {\n"
|
|
"${body}\n"
|
|
"};\n")
|
|
|
|
byte_strs = ['0x{:02X}'.format(b) for b in self._buffer.data]
|
|
|
|
lines = []
|
|
for byte_strs_for_line in utils.split_into_chunks(byte_strs, 12):
|
|
bytes_segment = ', '.join(byte_strs_for_line)
|
|
lines.append(f' {bytes_segment},')
|
|
|
|
return textwrap.indent(
|
|
buffer_template.substitute(buffer_name=self._buffer_name,
|
|
size=len(self._buffer.data),
|
|
body='\n'.join(lines)), indent)
|
|
|
|
|
|
class Tensor(object):
|
|
|
|
def __init__(self, buffer: Buffer, tensor: schema_fb.TensorT):
|
|
self._buffer = buffer
|
|
self._tensor: schema_fb.TensorT = tensor
|
|
|
|
@property
|
|
def buffer_index(self) -> bool:
|
|
return self._tensor.buffer
|
|
|
|
@property
|
|
def buffer(self) -> Buffer:
|
|
return self._buffer
|
|
|
|
@property
|
|
def has_shape(self) -> bool:
|
|
return self._tensor.shape is not None
|
|
|
|
@property
|
|
def needs_zero_length_int_array(self) -> bool:
|
|
return not self.has_shape
|
|
|
|
def generate_c_tensor_dims(self, type_name: str, tensor_name: str) -> str:
|
|
if not self.has_shape:
|
|
return f"// No data dims necessary for {tensor_name}"
|
|
return utils.IntArray(self._tensor.shape).generate_c_struct(
|
|
type_name + "Dims", tensor_name + "_dims")
|
|
|
|
def generate_c_tensor_init(self, tflite_tensor_name: str,
|
|
tensor_name: str) -> str:
|
|
init_template = string.Template(
|
|
"${tflite_tensor_name} = TfLiteEvalTensor{\n"
|
|
" .data = {.data = static_cast<void*>(${data})},\n"
|
|
" .dims = ${dims},\n"
|
|
" .type = ${tflite_type}};")
|
|
dims = "reinterpret_cast<TfLiteIntArray*>(&{})".format(
|
|
f"{tensor_name}_dims" if self._tensor.
|
|
shape is not None else "zero_length_int_array")
|
|
|
|
return init_template.substitute(
|
|
tflite_tensor_name=tflite_tensor_name,
|
|
tensor_name=tensor_name,
|
|
data=self._buffer.address,
|
|
dims=dims,
|
|
tflite_type=_TENSOR_TYPES[self._tensor.type])
|