Change swig dependency in pip_packge build to use pybind11.

PiperOrigin-RevId: 302850381
Change-Id: If7693293985c24d0e57c0129d1a73bc6e4ddc4a9
This commit is contained in:
Taehee Jeong 2020-03-25 02:56:42 -07:00 committed by TensorFlower Gardener
parent 0298dd6900
commit ccf558d3e2
2 changed files with 47 additions and 22 deletions

View File

@ -14,6 +14,7 @@ RUN apt-get update && \
python-setuptools \
python-wheel \
python-numpy \
python-pip \
libpython-dev \
libpython-dev:armhf \
libpython-dev:arm64 \
@ -21,6 +22,7 @@ RUN apt-get update && \
python3-setuptools \
python3-wheel \
python3-numpy \
python3-pip \
libpython3-dev \
libpython3-dev:armhf \
libpython3-dev:arm64 \
@ -29,8 +31,11 @@ RUN apt-get update && \
zlib1g-dev \
zlib1g-dev:armhf \
zlib1g-dev:arm64 \
swig \
curl \
unzip \
git && \
apt-get clean
RUN pip install pip --upgrade
RUN pip install pybind11
RUN pip3 install pip --upgrade
RUN pip3 install pybind11

View File

@ -24,9 +24,11 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import glob
import multiprocessing
import os
import subprocess
import sys
from distutils.command.build_ext import build_ext
import numpy
@ -65,7 +67,7 @@ for name in ['TARGET', 'TARGET_ARCH', 'CC_PREFIX', 'EXTRA_CXXFLAGS']:
# with more than 4GB, use all the CPUs, otherwise only 1.
def get_build_cpus():
physical_bytes = os.sysconf('SC_PAGESIZE') * os.sysconf('SC_PHYS_PAGES')
if physical_bytes < (1<<30) * 4:
if physical_bytes < (1 << 30) * 4:
return 1
else:
return multiprocessing.cpu_count()
@ -73,9 +75,9 @@ def get_build_cpus():
def make_args(target='', quiet=True):
"""Construct make command line."""
args = (['make', 'SHELL=/bin/bash',
'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR]
+ MAKE_CROSS_OPTIONS +
args = ([
'make', 'SHELL=/bin/bash', 'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR
] + MAKE_CROSS_OPTIONS +
['-f', RELATIVE_MAKEFILE_PATH, '-j',
str(get_build_cpus())])
if quiet:
@ -128,28 +130,46 @@ class CustomBuildPy(build_py, object):
return super(CustomBuildPy, self).run()
def get_pybind_include():
"""pybind11 include directory is not correctly resolved.
This fixes include directory to /usr/local/pythonX.X
Returns:
include directories to find pybind11
"""
if sys.version_info[0] == 3:
include_dirs = glob.glob('/usr/local/include/python3*')
else:
include_dirs = glob.glob('/usr/local/include/python2*')
for include_dir in include_dirs:
os.symlink(include_dir, os.path.join(include_dir, 'include'))
return include_dirs
LIB_TFLITE = 'tensorflow-lite'
LIB_TFLITE_DIR = make_output('libdir')
ext = Extension(
name='%s._interpreter_wrapper' % PACKAGE_NAME,
language='c++',
sources=['interpreter_wrapper/interpreter_wrapper.i',
'interpreter_wrapper/interpreter_wrapper.cc',
'interpreter_wrapper/numpy.cc',
'interpreter_wrapper/python_error_reporter.cc',
'interpreter_wrapper/python_utils.cc'],
sources=[
'interpreter_wrapper/interpreter_wrapper.cc',
'interpreter_wrapper/interpreter_wrapper_pybind11.cc',
'interpreter_wrapper/numpy.cc',
'interpreter_wrapper/python_error_reporter.cc',
'interpreter_wrapper/python_utils.cc'
],
extra_compile_args=['--std=c++11'],
swig_opts=['-c++',
'-I%s' % TENSORFLOW_DIR,
'-module', 'interpreter_wrapper',
'-outdir', PACKAGE_NAME],
include_dirs=[TENSORFLOW_DIR,
os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
'pip_package'),
numpy.get_include(),
os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
os.path.join(DOWNLOADS_DIR, 'absl')],
include_dirs=[
TENSORFLOW_DIR,
os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
'pip_package'),
numpy.get_include(),
os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
os.path.join(DOWNLOADS_DIR, 'absl')
] + get_pybind_include(),
libraries=[LIB_TFLITE],
library_dirs=[LIB_TFLITE_DIR])
@ -186,9 +206,9 @@ setup(
ext_modules=[ext],
install_requires=[
'numpy >= 1.16.0',
'pybind11 >= 2.4.3',
],
cmdclass={
'build_ext': CustomBuildExt,
'build_py': CustomBuildPy,
}
)
})