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 multiprocessing
28import os
29import subprocess
30
31from distutils.command.build_ext import build_ext
32import numpy
33
34from setuptools import Extension
35from setuptools import find_packages
36from setuptools import setup
37from setuptools.command.build_py import build_py
38PACKAGE_NAME = 'tflite-runtime'
39PACKAGE_VERSION = os.environ['TENSORFLOW_VERSION']
40DOCLINES = __doc__.split('\n')
41PACKAGE = 'tflite_runtime.lite.python'
42TENSORFLOW_DIR = os.environ['TENSORFLOW_SRC_ROOT']
43
44# Setup cross compiling
45TARGET = (
46    os.environ['TENSORFLOW_TARGET'] if 'TENSORFLOW_TARGET' in os.environ
47    else None)
48if TARGET == 'rpi':
49  os.environ['CXX'] = 'arm-linux-gnueabihf-g++'
50  os.environ['CC'] = 'arm-linux-gnueabihf-g++'
51MAKE_CROSS_OPTIONS = ['TARGET=%s' % TARGET]  if TARGET else []
52
53RELATIVE_MAKE_DIR = os.path.join('tensorflow', 'lite', 'tools', 'make')
54MAKE_DIR = os.path.join(TENSORFLOW_DIR, RELATIVE_MAKE_DIR)
55DOWNLOADS_DIR = os.path.join(MAKE_DIR, 'downloads')
56RELATIVE_MAKEFILE_PATH = os.path.join(RELATIVE_MAKE_DIR, 'Makefile')
57DOWNLOAD_SCRIPT_PATH = os.path.join(MAKE_DIR, 'download_dependencies.sh')
58
59
60# Check physical memory and if we are on a reasonable non small SOC machine
61# with more than 4GB, use all the CPUs, otherwisxe only 1.
62def get_build_cpus():
63  physical_bytes = os.sysconf('SC_PAGESIZE') * os.sysconf('SC_PHYS_PAGES')
64  if physical_bytes < (1<<30) * 4:
65    return 1
66  else:
67    return multiprocessing.cpu_count()
68
69
70def make_args(target='', quiet=True):
71  """Construct make command line."""
72  args = (['make', 'SHELL=/bin/bash', '-C', TENSORFLOW_DIR]
73          + MAKE_CROSS_OPTIONS +
74          ['-f', RELATIVE_MAKEFILE_PATH, '-j',
75           str(get_build_cpus())])
76  if quiet:
77    args.append('--quiet')
78  if target:
79    args.append(target)
80  return args
81
82
83def make_output(target):
84  """Invoke make on the target and return output."""
85  return subprocess.check_output(make_args(target)).decode('utf-8').strip()
86
87
88def make():
89  """Invoke make to build tflite C++ sources.
90
91  Build dependencies:
92     apt-get install swig libjpeg-dev zlib1g-dev python3-dev python3-nump
93  """
94  subprocess.check_call(make_args(quiet=False))
95
96
97def download_dependencies():
98  """Download build dependencies if haven't done yet."""
99  if not os.path.isdir(DOWNLOADS_DIR) or not os.listdir(DOWNLOADS_DIR):
100    subprocess.check_call(DOWNLOAD_SCRIPT_PATH)
101
102
103class CustomBuildExt(build_ext, object):
104
105  def run(self):
106    download_dependencies()
107    make()
108
109    return super(CustomBuildExt, self).run()
110
111
112class CustomBuildPy(build_py, object):
113
114  def run(self):
115    self.run_command('build_ext')
116    return super(CustomBuildPy, self).run()
117
118
119LIB_TFLITE = 'tensorflow-lite'
120LIB_TFLITE_DIR = make_output('libdir')
121
122ext = Extension(
123    name='%s._interpreter_wrapper' % PACKAGE,
124    language='c++',
125    sources=['interpreter_wrapper/interpreter_wrapper.i',
126             'interpreter_wrapper/interpreter_wrapper.cc'],
127    swig_opts=['-c++',
128               '-I%s' % TENSORFLOW_DIR,
129               '-module', 'interpreter_wrapper',
130               '-outdir', '.'],
131    extra_compile_args=['-std=c++11'],
132    include_dirs=[TENSORFLOW_DIR,
133                  os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
134                               'pip_package'),
135                  numpy.get_include(),
136                  os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
137                  os.path.join(DOWNLOADS_DIR, 'absl')],
138    libraries=[LIB_TFLITE],
139    library_dirs=[LIB_TFLITE_DIR])
140
141
142setup(
143    name=PACKAGE_NAME,
144    version=PACKAGE_VERSION,
145    description=DOCLINES[0],
146    long_description='\n'.join(DOCLINES[2:]),
147    url='https://www.tensorflow.org/lite/',
148    author='Google Inc.',
149    author_email='packages@tensorflow.org',
150    license='Apache 2.0',
151    include_package_data=True,
152    keywords='tflite tensorflow tensor machine learning',
153    packages=find_packages(exclude=[]),
154    ext_modules=[ext],
155    package_dir={PACKAGE: '.'},
156    cmdclass={
157        'build_ext': CustomBuildExt,
158        'build_py': CustomBuildPy,
159    }
160)
161