support tf installation and separate it from pytorch (intel-analytics…

…#4535)
Le-Zheng · Aug 23, 2021 · 534522a · 534522a
1 parent ed12137
commit 534522a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 10 deletions.
diff --git a/python/nano/script/bigdl-nano-init b/python/nano/script/bigdl-nano-init
@@ -91,7 +91,13 @@ OPENMP=0
 JEMALLOC=0
 
 # Find conda dir
-BASE_DIR="$(dirname "$0")/../"
+if [ $0 != $BASH_SOURCE ]; then
+        # this script is being sourced
+        BASE_DIR="$(dirname "$BASH_SOURCE")/.."
+else
+        # this script is being executed
+        BASE_DIR="$(dirname "$0")/../"
+fi
 echo "conda dir found: $BASE_DIR"
 LIB_DIR=$BASE_DIR/lib
 

diff --git a/python/nano/setup.py b/python/nano/setup.py
@@ -104,14 +104,18 @@ def setup_package():
                            f"-cp{py_version.major}{py_version.minor}m-linux_x86_64.whl"
     torchvision_link = "https://download.pytorch.org/whl/" + torchvision_links[torchvision_whl_name]
 
-    install_requires_list = ["pytorch_lightning",
-                             "opencv-python-headless",
-                             "PyTurboJPEG",
-                             "opencv-transforms",
-                             "intel-openmp",
-                             "torchvision",
-                             f"torch_ipex @ {ipex_link}",
-                             f"torchvision @ {torchvision_link}"]
+    install_requires = ["intel-openmp"]
+
+    tensorflow_requires = ["intel-tensorflow"]
+
+    pytorch_requires = ["pytorch_lightning",
+                        "opencv-python-headless",
+                        "PyTurboJPEG",
+                        "opencv-transforms",
+                        "torchvision",
+                        f"torch_ipex @ {ipex_link}",
+                        f"torchvision @ {torchvision_link}"]
+
 
     download_libs()
 
@@ -122,7 +126,8 @@ def setup_package():
         author='',
         author_email='',
         url='https://github.com/intel-analytics/analytics-zoo/tree/bigdl-2.0',
-        install_requires=install_requires_list,
+        install_requires=install_requires,
+        extras_require={"tensorflow": tensorflow_requires, "pytorch": pytorch_requires},
         packages=get_nano_packages(),
         package_data={"bigdl.nano": ["libs/libjemalloc.so"]},
         package_dir={'': 'src'},