-
-
Notifications
You must be signed in to change notification settings - Fork 3
/
setup.py
34 lines (30 loc) · 821 Bytes
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from setuptools import setup, find_packages
import platform
import os
def check_tensorrt():
try:
import tensorrt
return True
except:
return False
def check_jetson():
if platform.machine() == "aarch64":
return True
else:
return False
def get_requires():
requires = []
if not check_jetson():
requires = ["tensorflow==2.11.0", "tensorflow_addons"]
requires += ["torch", "tensorflow==2.11.0", "tensorflow_probability", "onnx", "onnx_tf @ git+https://github.com/onnx/onnx-tensorflow",
"matplotlib", "onnx-simplifier"]
if check_tensorrt():
requires += ["pycuda"]
return requires
setup(
name="nne",
scripts=["bin/nne"],
packages=find_packages(),
install_requires=get_requires(),
version="0.1"
)