address review commments
diff --git a/tensorflow/lite/examples/python/label_image.py b/tensorflow/lite/examples/python/label_image.py
index c6e0fbd..9d3c47f 100644
--- a/tensorflow/lite/examples/python/label_image.py
+++ b/tensorflow/lite/examples/python/label_image.py
@@ -57,13 +57,13 @@
help='input standard deviation')
parser.add_argument(
'--num_threads',
- default=1,
+ default=1, type=int,
help='number of threads')
args = parser.parse_args()
interpreter = tf.lite.Interpreter(
model_path=args.model_file,
- num_threads=int(args.num_threads))
+ num_threads=args.num_threads)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
@@ -100,4 +100,5 @@
else:
print('{:08.6f}: {}'.format(float(results[i] / 255.0), labels[i]))
- print("time: ", stop_time - start_time)
+ #print("time: ", stop_time - start_time)
+ print('time: {:.3f}ms'.format((stop_time - start_time) * 1000))
diff --git a/tensorflow/lite/python/interpreter.py b/tensorflow/lite/python/interpreter.py
index 4c25280..36ea2e2 100644
--- a/tensorflow/lite/python/interpreter.py
+++ b/tensorflow/lite/python/interpreter.py
@@ -523,6 +523,7 @@
def reset_all_variables(self):
return self._interpreter.ResetVariableTensors()
+
class InterpreterWithCustomOps(Interpreter):
"""Interpreter interface for TensorFlow Lite Models that accepts custom ops.
diff --git a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc
index c457e68..2a8c1ff 100644
--- a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc
+++ b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc
@@ -706,9 +706,9 @@
Py_RETURN_NONE;
}
-PyObject* InterpreterWrapper::SetNumThreads(int i) {
+PyObject* InterpreterWrapper::SetNumThreads(int num_threads) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
- interpreter_->SetNumThreads(i);
+ interpreter_->SetNumThreads(num_threads);
Py_RETURN_NONE;
}
diff --git a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h
index d714118..b799a30 100644
--- a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h
+++ b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h
@@ -87,7 +87,7 @@
// should be the interpreter object providing the memory.
PyObject* tensor(PyObject* base_object, int i);
- PyObject* SetNumThreads(int i);
+ PyObject* SetNumThreads(int num_threads);
// Adds a delegate to the interpreter.
PyObject* ModifyGraphWithDelegate(TfLiteDelegate* delegate);
diff --git a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
index 55c377c..74bbf6f 100644
--- a/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
+++ b/tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
@@ -149,7 +149,7 @@
.def(
"SetNumThreads",
[](InterpreterWrapper& self, int i) {
- return tensorflow::pyo_or_throw(self.SetNumThreads(i));
+ return tensorflow::PyoOrThrow(self.SetNumThreads(i));
},
R"pbdoc(
ask the interpreter to set the number of threads to use.
