MLBEDSW-2067: added custom exceptions
Added custom exceptions to handle different types of input errors.
Also performed minor formatting changes using flake8/black.
Change-Id: Ie5b05361507d5e569aff045757aec0a4a755ae98
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
diff --git a/ethosu/vela/model_reader.py b/ethosu/vela/model_reader.py
index d1cdc9b..6deb253 100644
--- a/ethosu/vela/model_reader.py
+++ b/ethosu/vela/model_reader.py
@@ -15,6 +15,9 @@
# limitations under the License.
# Description:
# Dispatcher for reading a neural network model.
+from . import tflite_reader
+from .errors import InputFileError
+from .errors import VelaError
class ModelReaderOptions:
@@ -29,15 +32,17 @@
def read_model(fname, options, feed_dict={}, output_node_names=[], initialisation_nodes=[]):
if fname.endswith(".tflite"):
- from . import tflite_reader
-
- nng = tflite_reader.read_tflite(
- fname,
- options.batch_size,
- feed_dict=feed_dict,
- output_node_names=output_node_names,
- initialisation_nodes=initialisation_nodes,
- )
+ try:
+ return tflite_reader.read_tflite(
+ fname,
+ options.batch_size,
+ feed_dict=feed_dict,
+ output_node_names=output_node_names,
+ initialisation_nodes=initialisation_nodes,
+ )
+ except VelaError as e:
+ raise e
+ except Exception as e:
+ raise InputFileError(fname, str(e))
else:
- assert 0, "Unknown model format"
- return nng
+ raise InputFileError(fname, "Unknown input file format. Only .tflite files are supported")