Move TOCO python bindings to be built under the main TensorFlow Python DSO.

This allows future linking against MLIR without bloating binary. This
was not done initially since TF lite was in contrib.

PiperOrigin-RevId: 257895424
This commit is contained in:
Andrew Selle 2019-07-12 17:02:17 -07:00 committed by TensorFlower Gardener
parent 5e1c78311d
commit f1a2ce4406
13 changed files with 224 additions and 182 deletions

View File

@ -174,7 +174,7 @@ py_library(
"wrap_toco.py",
],
deps = [
"//tensorflow/lite/toco/python:tensorflow_wrap_toco",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python:util",
],
)

View File

@ -10,6 +10,8 @@ cc_library(
srcs = ["numpy.cc"],
hdrs = ["numpy.h"],
deps = [
"//tensorflow/lite:string_util",
"//tensorflow/lite/c:c_api_internal",
"//third_party/py/numpy:headers",
"//third_party/python_runtime:headers",
],
@ -27,7 +29,6 @@ cc_library(
"//tensorflow/lite:string_util",
"//tensorflow/lite/c:c_api_internal",
"//tensorflow/lite/kernels:builtin_ops",
"//third_party/py/numpy:headers",
"//third_party/python_runtime:headers",
"@com_google_absl//absl/memory",
],
@ -48,10 +49,8 @@ cc_library(
srcs = ["python_utils.cc"],
hdrs = ["python_utils.h"],
deps = [
":numpy",
"//tensorflow/lite:framework",
"//tensorflow/lite:string_util",
"//third_party/py/numpy:headers",
"//third_party/python_runtime:headers",
],
)

View File

@ -16,10 +16,159 @@ limitations under the License.
#define TFLITE_IMPORT_NUMPY // See numpy.h for explanation.
#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
#include <memory>
namespace tflite {
namespace python {
void ImportNumpy() { import_array1(); }
} // namespace python
namespace python_utils {
struct PyObjectDereferencer {
void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
};
using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
switch (tf_lite_type) {
case kTfLiteFloat32:
return NPY_FLOAT32;
case kTfLiteFloat16:
return NPY_FLOAT16;
case kTfLiteInt32:
return NPY_INT32;
case kTfLiteInt16:
return NPY_INT16;
case kTfLiteUInt8:
return NPY_UINT8;
case kTfLiteInt8:
return NPY_INT8;
case kTfLiteInt64:
return NPY_INT64;
case kTfLiteString:
return NPY_STRING;
case kTfLiteBool:
return NPY_BOOL;
case kTfLiteComplex64:
return NPY_COMPLEX64;
case kTfLiteNoType:
return NPY_NOTYPE;
// Avoid default so compiler errors created when new types are made.
}
return NPY_NOTYPE;
}
TfLiteType TfLiteTypeFromPyType(int py_type) {
switch (py_type) {
case NPY_FLOAT32:
return kTfLiteFloat32;
case NPY_INT32:
return kTfLiteInt32;
case NPY_INT16:
return kTfLiteInt16;
case NPY_UINT8:
return kTfLiteUInt8;
case NPY_INT8:
return kTfLiteInt8;
case NPY_INT64:
return kTfLiteInt64;
case NPY_BOOL:
return kTfLiteBool;
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE:
return kTfLiteString;
case NPY_COMPLEX64:
return kTfLiteComplex64;
// Avoid default so compiler errors created when new types are made.
}
return kTfLiteNoType;
}
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
int pyarray_type = PyArray_TYPE(array);
return TfLiteTypeFromPyType(pyarray_type);
}
#if PY_VERSION_HEX >= 0x03030000
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
Py_ssize_t len = -1;
const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
if (buf == NULL) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#else
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
if (!utemp) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
return false;
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#endif
bool FillStringBufferFromPyString(PyObject* value,
DynamicBuffer* dynamic_buffer) {
if (PyUnicode_Check(value)) {
return FillStringBufferFromPyUnicode(value, dynamic_buffer);
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer) {
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
switch (PyArray_TYPE(array)) {
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE: {
UniquePyObjectRef iter(PyArray_IterNew(value));
while (PyArray_ITER_NOTDONE(iter.get())) {
UniquePyObjectRef item(PyArray_GETITEM(
array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
return false;
}
PyArray_ITER_NEXT(iter.get());
}
return true;
}
default:
break;
}
PyErr_Format(PyExc_ValueError,
"Cannot use numpy array of type %d for string tensor.",
PyArray_TYPE(array));
return false;
}
} // namespace python_utils
} // namespace tflite

View File

@ -50,6 +50,8 @@ limitations under the License.
#include "numpy/arrayobject.h"
#include "numpy/ufuncobject.h"
#include "tensorflow/lite/c/c_api_internal.h"
#include "tensorflow/lite/string_util.h"
namespace tflite {
namespace python {
@ -57,6 +59,19 @@ namespace python {
void ImportNumpy();
} // namespace python
namespace python_utils {
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
TfLiteType TfLiteTypeFromPyType(int py_type);
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer);
} // namespace python_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_

View File

@ -17,155 +17,9 @@ limitations under the License.
#include <memory>
#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
namespace tflite {
namespace python_utils {
struct PyObjectDereferencer {
void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
};
using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
switch (tf_lite_type) {
case kTfLiteFloat32:
return NPY_FLOAT32;
case kTfLiteFloat16:
return NPY_FLOAT16;
case kTfLiteInt32:
return NPY_INT32;
case kTfLiteInt16:
return NPY_INT16;
case kTfLiteUInt8:
return NPY_UINT8;
case kTfLiteInt8:
return NPY_INT8;
case kTfLiteInt64:
return NPY_INT64;
case kTfLiteString:
return NPY_STRING;
case kTfLiteBool:
return NPY_BOOL;
case kTfLiteComplex64:
return NPY_COMPLEX64;
case kTfLiteNoType:
return NPY_NOTYPE;
// Avoid default so compiler errors created when new types are made.
}
return NPY_NOTYPE;
}
TfLiteType TfLiteTypeFromPyType(int py_type) {
switch (py_type) {
case NPY_FLOAT32:
return kTfLiteFloat32;
case NPY_INT32:
return kTfLiteInt32;
case NPY_INT16:
return kTfLiteInt16;
case NPY_UINT8:
return kTfLiteUInt8;
case NPY_INT8:
return kTfLiteInt8;
case NPY_INT64:
return kTfLiteInt64;
case NPY_BOOL:
return kTfLiteBool;
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE:
return kTfLiteString;
case NPY_COMPLEX64:
return kTfLiteComplex64;
// Avoid default so compiler errors created when new types are made.
}
return kTfLiteNoType;
}
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
int pyarray_type = PyArray_TYPE(array);
return TfLiteTypeFromPyType(pyarray_type);
}
#if PY_VERSION_HEX >= 0x03030000
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
Py_ssize_t len = -1;
const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
if (buf == NULL) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#else
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
if (!utemp) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
return false;
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#endif
bool FillStringBufferFromPyString(PyObject* value,
DynamicBuffer* dynamic_buffer) {
if (PyUnicode_Check(value)) {
return FillStringBufferFromPyUnicode(value, dynamic_buffer);
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer) {
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
switch (PyArray_TYPE(array)) {
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE: {
UniquePyObjectRef iter(PyArray_IterNew(value));
while (PyArray_ITER_NOTDONE(iter.get())) {
UniquePyObjectRef item(PyArray_GETITEM(
array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
return false;
}
PyArray_ITER_NEXT(iter.get());
}
return true;
}
default:
break;
}
PyErr_Format(PyExc_ValueError,
"Cannot use numpy array of type %d for string tensor.",
PyArray_TYPE(array));
return false;
}
int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length) {
#if PY_MAJOR_VERSION >= 3
return PyBytes_AsStringAndSize(obj, data, length);

View File

@ -16,8 +16,9 @@ limitations under the License.
#ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
#include <Python.h>
#include "tensorflow/lite/context.h"
#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
#include "tensorflow/lite/string_util.h"
namespace tflite {
@ -27,15 +28,6 @@ struct PyDecrefDeleter {
void operator()(PyObject* p) const { Py_DECREF(p); }
};
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
TfLiteType TfLiteTypeFromPyType(int py_type);
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer);
int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length);
PyObject* ConvertToPyString(const char* data, size_t length);

View File

@ -17,22 +17,15 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensorflow.python.util.lazy_loader import LazyLoader
from tensorflow.python import pywrap_tensorflow
# TODO(b/131123224): Lazy load since some of the performance benchmark skylark
# rules and monolithic build break dependencies.
_toco_python = LazyLoader(
"tensorflow_wrap_toco", globals(),
"tensorflow.lite.toco.python."
"tensorflow_wrap_toco")
del LazyLoader
# TODO(b/137402359): Remove lazy loading wrapper
def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str,
debug_info_str, enable_mlir_converter):
"""Wraps TocoConvert with lazy loader."""
return _toco_python.TocoConvert(
return pywrap_tensorflow.TocoConvert(
model_flags_str,
toco_flags_str,
input_data_str,
@ -43,4 +36,4 @@ def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str,
def wrapped_get_potentially_supported_ops():
"""Wraps TocoGetPotentiallySupportedOps with lazy loader."""
return _toco_python.TocoGetPotentiallySupportedOps()
return pywrap_tensorflow.TocoGetPotentiallySupportedOps()

View File

@ -22,7 +22,9 @@ cc_library(
name = "toco_python_api",
srcs = ["toco_python_api.cc"],
hdrs = ["toco_python_api.h"],
features = ["no_layering_check"],
visibility = [
"//tensorflow/python:__subpackages__",
],
deps = [
"//third_party/python_runtime:headers",
"//tensorflow/core:lib",
@ -43,20 +45,17 @@ cc_library(
}),
)
tf_py_wrap_cc(
# Compatibility stub. Remove when internal customers moved.
py_library(
name = "tensorflow_wrap_toco",
srcs = ["toco.i"],
srcs = ["tensorflow_wrap_toco.py"],
visibility = [
"//learning/expander/pod/deep_pod/utils:__subpackages__",
"//research/handwriting/converters/tflite:__subpackages__",
"//tensorflow/lite:__subpackages__",
],
deps = [
":toco_python_api",
"//tensorflow/lite/toco:model_flags_proto_cc",
"//tensorflow/lite/toco:toco_flags_proto_cc",
"//third_party/python_runtime:headers",
"@com_google_absl//absl/strings",
"//tensorflow/python:pywrap_tensorflow",
],
)
@ -66,8 +65,8 @@ py_binary(
python_version = "PY2",
srcs_version = "PY2AND3",
deps = [
":tensorflow_wrap_toco",
"//tensorflow/python:platform",
"//tensorflow/python:pywrap_tensorflow",
],
)
@ -84,3 +83,10 @@ tf_py_test(
"no_pip",
],
)
exports_files(
["toco.i"],
visibility = [
"//tensorflow/python:__subpackages__",
],
)

View File

@ -0,0 +1,24 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Stub to make toco convert accessible."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# TODO(aselle): Remove once no clients internally depend on this.
# pylint: disable=unused-import
from tensorflow.python.pywrap_tensorflow import TocoConvert
# pylint: enable=unused-import

View File

@ -19,6 +19,12 @@ limitations under the License.
#include "tensorflow/lite/toco/python/toco_python_api.h"
%}
// The TensorFlow exception handler releases the GIL with
// Py_BEGIN_ALLOW_THREADS. Remove that because these function use the Python
// API to decode inputs.
%noexception toco::TocoConvert;
%noexception toco::TocoGetPotentiallySupportedOps;
namespace toco {
// Convert a model represented in `input_contents`. `model_flags_proto`

View File

@ -19,7 +19,7 @@ from __future__ import print_function
import argparse
import sys
from tensorflow.lite.toco.python import tensorflow_wrap_toco
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.platform import app
FLAGS = None
@ -43,7 +43,7 @@ def execute(unused_args):
enable_mlir_converter = FLAGS.enable_mlir_converter
output_str = tensorflow_wrap_toco.TocoConvert(
output_str = pywrap_tensorflow.TocoConvert(
model_str,
toco_str,
input_str,

View File

@ -4759,6 +4759,7 @@ tf_py_wrap_cc(
"util/tfprof.i",
"util/transform_graph.i",
"util/util.i",
"//tensorflow/lite/toco/python:toco.i",
],
# add win_def_file for pywrap_tensorflow
win_def_file = select({
@ -4805,6 +4806,7 @@ tf_py_wrap_cc(
"//tensorflow/core/distributed_runtime:server_lib",
"//tensorflow/core/profiler/internal:print_model_analysis",
"//tensorflow/tools/graph_transforms:transform_graph_lib",
"//tensorflow/lite/toco/python:toco_python_api",
"//tensorflow/python/eager:pywrap_tfe_lib",
] + (tf_additional_lib_deps() +
tf_additional_plugin_deps() +

View File

@ -36,6 +36,8 @@ limitations under the License.
%include "tensorflow/python/lib/core/bfloat16.i"
%include "tensorflow/lite/toco/python/toco.i"
%include "tensorflow/python/lib/io/file_io.i"
%include "tensorflow/python/training/quantize_training.i"