blob: 6deb2538b8d41539faa4e637cbf360208a9db210 [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# Dispatcher for reading a neural network model.
Louis Verhaard7db78962020-05-25 15:05:26 +020018from . import tflite_reader
19from .errors import InputFileError
20from .errors import VelaError
Tim Hall79d07d22020-04-27 18:20:16 +010021
22
23class ModelReaderOptions:
24 def __init__(self, batch_size=1):
25 self.batch_size = batch_size
26
27 def __str__(self):
28 return type(self).__name__ + ": " + str(self.__dict__)
29
30 __repr__ = __str__
31
32
33def read_model(fname, options, feed_dict={}, output_node_names=[], initialisation_nodes=[]):
34 if fname.endswith(".tflite"):
Louis Verhaard7db78962020-05-25 15:05:26 +020035 try:
36 return tflite_reader.read_tflite(
37 fname,
38 options.batch_size,
39 feed_dict=feed_dict,
40 output_node_names=output_node_names,
41 initialisation_nodes=initialisation_nodes,
42 )
43 except VelaError as e:
44 raise e
45 except Exception as e:
46 raise InputFileError(fname, str(e))
Tim Hall79d07d22020-04-27 18:20:16 +010047 else:
Louis Verhaard7db78962020-05-25 15:05:26 +020048 raise InputFileError(fname, "Unknown input file format. Only .tflite files are supported")