Make TensorFlow Lite Interpreter release the GIL to enable multithreading.

PiperOrigin-RevId: 240187848
This commit is contained in:
Andrew Selle 2019-03-25 11:53:01 -07:00 committed by TensorFlower Gardener
parent 5911c8a43e
commit 97bebc917b
2 changed files with 25 additions and 3 deletions

View File

@ -49,8 +49,18 @@ except ImportError:
@_tf_export('lite.Interpreter') @_tf_export('lite.Interpreter')
class Interpreter(object): class Interpreter(object):
"""Interpreter interface for TF-Lite Models.""" """Interpreter interface for TensorFlow Lite Models.
This makes the TensorFlow Lite interpreter accessible in Python.
It is possible to use this interpreter in a multithreaded Python environment,
but you must be sure to call functions of a particular instance from only
one thread at a time. So if you want to have 4 threads running different
inferences simultaneously, create an interpreter for each one as thread-local
data. Similarly, if you are calling invoke() in one thread on a single
interpreter but you want to use tensor() on another thread once it is done,
you must use a synchronization primitive between the threads to ensure invoke
has returned before calling tensor().
"""
def __init__(self, model_path=None, model_content=None): def __init__(self, model_path=None, model_content=None):
"""Constructor. """Constructor.
@ -282,7 +292,10 @@ class Interpreter(object):
"""Invoke the interpreter. """Invoke the interpreter.
Be sure to set the input sizes, allocate tensors and fill values before Be sure to set the input sizes, allocate tensors and fill values before
calling this. calling this. Also, note that this function releases the GIL so heavy
computation can be done in the background while the Python interpreter
continues. No other function on this object should be called while the
invoke() call has not finished.
Raises: Raises:
ValueError: When the underlying interpreter fails raise ValueError. ValueError: When the underlying interpreter fails raise ValueError.

View File

@ -125,7 +125,16 @@ PyObject* InterpreterWrapper::AllocateTensors() {
PyObject* InterpreterWrapper::Invoke() { PyObject* InterpreterWrapper::Invoke() {
TFLITE_PY_ENSURE_VALID_INTERPRETER(); TFLITE_PY_ENSURE_VALID_INTERPRETER();
TFLITE_PY_CHECK(interpreter_->Invoke());
// Release the GIL so that we can run multiple interpreters in parallel
TfLiteStatus status_code = kTfLiteOk;
Py_BEGIN_ALLOW_THREADS; // To return can happen between this and end!
status_code = interpreter_->Invoke();
Py_END_ALLOW_THREADS;
TFLITE_PY_CHECK(
status_code); // don't move this into the Py_BEGIN/Py_End block
Py_RETURN_NONE; Py_RETURN_NONE;
} }