• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""TensorFlow Lite is for mobile and embedded devices.
16
17TensorFlow Lite is the official solution for running machine learning models on
18mobile and embedded devices. It enables on-device machine learning inference
19with low latency and a small binary size on Android, iOS, and other operating
20systems.
21"""
22
23from __future__ import absolute_import
24from __future__ import division
25from __future__ import print_function
26
27import glob
28import multiprocessing
29import os
30import subprocess
31import sys
32import sysconfig
33
34from distutils.command.build_ext import build_ext
35import numpy
36import pybind11
37
38from setuptools import Extension
39from setuptools import find_packages
40from setuptools import setup
41from setuptools.command.build_py import build_py
42PACKAGE_NAME = 'tflite_runtime'
43PACKAGE_VERSION = os.environ['PACKAGE_VERSION']
44DOCLINES = __doc__.split('\n')
45TENSORFLOW_DIR = os.environ['TENSORFLOW_DIR']
46RELATIVE_MAKE_DIR = os.path.join('tensorflow', 'lite', 'tools', 'make')
47MAKE_DIR = os.path.join(TENSORFLOW_DIR, RELATIVE_MAKE_DIR)
48DOWNLOADS_DIR = os.path.join(MAKE_DIR, 'downloads')
49RELATIVE_MAKEFILE_PATH = os.path.join(RELATIVE_MAKE_DIR, 'Makefile')
50DOWNLOAD_SCRIPT_PATH = os.path.join(MAKE_DIR, 'download_dependencies.sh')
51
52# Setup cross compiling
53TARGET = os.environ.get('TENSORFLOW_TARGET')
54if TARGET == 'rpi':
55  os.environ['CXX'] = 'arm-linux-gnueabihf-g++'
56  os.environ['CC'] = 'arm-linux-gnueabihf-gcc'
57elif TARGET == 'aarch64':
58  os.environ['CXX'] = 'aarch64-linux-gnu-g++'
59  os.environ['CC'] = 'aarch64-linux-gnu-gcc'
60
61MAKE_CROSS_OPTIONS = []
62for name in [
63    'TARGET', 'TARGET_ARCH', 'CC_PREFIX', 'EXTRA_CXXFLAGS', 'EXTRA_CFLAGS'
64]:
65  value = os.environ.get('TENSORFLOW_%s' % name)
66  if value:
67    MAKE_CROSS_OPTIONS.append('%s=%s' % (name, value))
68
69
70# Check physical memory and if we are on a reasonable non small SOC machine
71# with more than 4GB, use all the CPUs, otherwise only 1.
72def get_build_cpus():
73  physical_bytes = os.sysconf('SC_PAGESIZE') * os.sysconf('SC_PHYS_PAGES')
74  if physical_bytes < (1 << 30) * 4:
75    return 1
76  else:
77    return multiprocessing.cpu_count()
78
79
80def make_args(target='', quiet=True):
81  """Construct make command line."""
82  args = ([
83      'make', 'SHELL=/bin/bash', 'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR
84  ] + MAKE_CROSS_OPTIONS +
85          ['-f', RELATIVE_MAKEFILE_PATH, '-j',
86           str(get_build_cpus())])
87  if quiet:
88    args.append('--quiet')
89  if target:
90    args.append(target)
91  return args
92
93
94def make_output(target):
95  """Invoke make on the target and return output."""
96  return subprocess.check_output(make_args(target)).decode('utf-8').strip()
97
98
99def make():
100  """Invoke make to build tflite C++ sources.
101
102  Build dependencies:
103     apt-get install swig libjpeg-dev zlib1g-dev python3-dev python3-nump
104  """
105  subprocess.check_call(make_args(quiet=False))
106
107
108def download_dependencies():
109  """Download build dependencies if haven't done yet."""
110  if not os.path.isdir(DOWNLOADS_DIR) or not os.listdir(DOWNLOADS_DIR):
111    subprocess.check_call(DOWNLOAD_SCRIPT_PATH)
112
113
114class CustomBuildExt(build_ext, object):
115  """Customized build extension."""
116
117  def get_ext_filename(self, ext_name):
118    if TARGET:
119      ext_path = ext_name.split('.')
120      return os.path.join(*ext_path) + '.so'
121    return super(CustomBuildExt, self).get_ext_filename(ext_name)
122
123  def run(self):
124    download_dependencies()
125    make()
126
127    return super(CustomBuildExt, self).run()
128
129
130class CustomBuildPy(build_py, object):
131
132  def run(self):
133    self.run_command('build_ext')
134    return super(CustomBuildPy, self).run()
135
136
137def get_pybind_include():
138  """pybind11 include directory is not correctly resolved.
139
140  This fixes include directory to /usr/local/pythonX.X
141
142  Returns:
143    include directories to find pybind11
144  """
145  if sys.version_info[0] == 3:
146    include_dirs = glob.glob('/usr/local/include/python3*')
147  else:
148    include_dirs = glob.glob('/usr/local/include/python2*')
149  include_dirs.append(sysconfig.get_path('include'))
150  tmp_include_dirs = []
151  pip_dir = os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
152                         'pip_package', 'gen')
153  for include_dir in include_dirs:
154    tmp_include_dir = os.path.join(pip_dir, include_dir[1:])
155    tmp_include_dirs.append(tmp_include_dir)
156    try:
157      os.makedirs(tmp_include_dir)
158      os.symlink(include_dir, os.path.join(tmp_include_dir, 'include'))
159    except IOError:  # file already exists.
160      pass
161  return tmp_include_dirs
162
163
164LIB_TFLITE = 'tensorflow-lite'
165LIB_TFLITE_DIR = make_output('libdir')
166
167ext = Extension(
168    name='%s._pywrap_tensorflow_interpreter_wrapper' % PACKAGE_NAME,
169    language='c++',
170    sources=[
171        'interpreter_wrapper/interpreter_wrapper.cc',
172        'interpreter_wrapper/interpreter_wrapper_pybind11.cc',
173        'interpreter_wrapper/numpy.cc',
174        'interpreter_wrapper/python_error_reporter.cc',
175        'interpreter_wrapper/python_utils.cc'
176    ],
177    extra_compile_args=['--std=c++11'],
178    include_dirs=[
179        TENSORFLOW_DIR,
180        os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
181                     'pip_package'),
182        numpy.get_include(),
183        os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
184        os.path.join(DOWNLOADS_DIR, 'absl'),
185        pybind11.get_include()
186    ],
187    libraries=[LIB_TFLITE],
188    library_dirs=[LIB_TFLITE_DIR])
189
190setup(
191    name=PACKAGE_NAME.replace('_', '-'),
192    version=PACKAGE_VERSION,
193    description=DOCLINES[0],
194    long_description='\n'.join(DOCLINES[2:]),
195    url='https://www.tensorflow.org/lite/',
196    author='Google, LLC',
197    author_email='packages@tensorflow.org',
198    license='Apache 2.0',
199    include_package_data=True,
200    keywords='tflite tensorflow tensor machine learning',
201    classifiers=[
202        'Development Status :: 5 - Production/Stable',
203        'Intended Audience :: Developers',
204        'Intended Audience :: Education',
205        'Intended Audience :: Science/Research',
206        'License :: OSI Approved :: Apache Software License',
207        'Programming Language :: Python :: 3',
208        'Programming Language :: Python :: 3.4',
209        'Programming Language :: Python :: 3.5',
210        'Programming Language :: Python :: 3.6',
211        'Programming Language :: Python :: 3.7',
212        'Programming Language :: Python :: 3.8',
213        'Topic :: Scientific/Engineering',
214        'Topic :: Scientific/Engineering :: Mathematics',
215        'Topic :: Scientific/Engineering :: Artificial Intelligence',
216        'Topic :: Software Development',
217        'Topic :: Software Development :: Libraries',
218        'Topic :: Software Development :: Libraries :: Python Modules',
219    ],
220    packages=find_packages(exclude=[]),
221    ext_modules=[ext],
222    install_requires=[
223        'numpy >= 1.16.0',
224    ],
225    cmdclass={
226        'build_ext': CustomBuildExt,
227        'build_py': CustomBuildPy,
228    })
229