blob: 0f79f9b2ab156d9e52bf5710046b39de8aa8d632 [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
Tim Hall79d07d22020-04-27 18:20:16 +010020
21
22class ModelReaderOptions:
23 def __init__(self, batch_size=1):
24 self.batch_size = batch_size
25
26 def __str__(self):
27 return type(self).__name__ + ": " + str(self.__dict__)
28
29 __repr__ = __str__
30
31
32def read_model(fname, options, feed_dict={}, output_node_names=[], initialisation_nodes=[]):
33 if fname.endswith(".tflite"):
Tim Hallc8310b12020-06-17 14:53:11 +010034 return tflite_reader.read_tflite(
35 fname,
36 options.batch_size,
37 feed_dict=feed_dict,
38 output_node_names=output_node_names,
39 initialisation_nodes=initialisation_nodes,
40 )
Tim Hall79d07d22020-04-27 18:20:16 +010041 else:
Tim Hallc8310b12020-06-17 14:53:11 +010042 raise InputFileError(fname, "Unsupported file extension. Only .tflite files are supported")