69 lines
2.3 KiB
Python
69 lines
2.3 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.
|
|
# ==============================================================================
|
|
""" Generates C/C++ inference source code. """
|
|
|
|
import pathlib
|
|
|
|
from mako import template
|
|
from typing import TypedDict
|
|
|
|
from tflite_micro.codegen import graph
|
|
|
|
_TEMPLATE_DIR = pathlib.Path(__file__).parent / 'templates'
|
|
_HEADER_TEMPLATE = _TEMPLATE_DIR / 'inference.h.mako'
|
|
_SOURCE_TEMPLATE = _TEMPLATE_DIR / 'inference.cc.mako'
|
|
|
|
|
|
class ModelData(TypedDict):
|
|
header_file: str
|
|
model_name: str
|
|
op_code_table: graph.OpCodeTable
|
|
graph: graph.Graph
|
|
|
|
|
|
def _render(output_file: pathlib.Path, template_file: pathlib.Path,
|
|
model_data: ModelData) -> None:
|
|
print("Generating {}".format(output_file))
|
|
t = template.Template(filename=str(template_file))
|
|
with output_file.open('w+') as file:
|
|
file.write(t.render(**model_data))
|
|
|
|
|
|
def _generate_header(header_path: pathlib.Path, model_data: ModelData) -> None:
|
|
_render(header_path, _HEADER_TEMPLATE, model_data)
|
|
|
|
|
|
def _generate_source(source_path: pathlib.Path, model_data: ModelData) -> None:
|
|
_render(source_path, _SOURCE_TEMPLATE, model_data)
|
|
|
|
|
|
def generate(output_dir: str, output_name: str,
|
|
op_code_table: graph.OpCodeTable, graph: graph.Graph) -> None:
|
|
""" Generate C/C++ inference code. """
|
|
header_file = f"{output_name}.h"
|
|
model_data: ModelData = {
|
|
'header_file': header_file,
|
|
'model_name': output_name,
|
|
'op_code_table': op_code_table,
|
|
'graph': graph,
|
|
}
|
|
|
|
# Ensure output directory exists
|
|
output_path = pathlib.Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
_generate_header(output_path / header_file, model_data)
|
|
_generate_source(output_path / f"{output_name}.cc", model_data)
|