├── .gitignore ├── .gitmodules ├── .travis.yml ├── CMakeLists.txt ├── Jenkinsfile ├── LICENSE ├── Makefile ├── NEWS.md ├── README.md ├── amalgamation ├── .gitignore ├── Makefile ├── README ├── amalgamation.py └── generate.py ├── cmake └── Utils.cmake ├── docs ├── .gitignore ├── Doxyfile ├── Makefile ├── README.txt ├── api │ └── python │ │ ├── compiler.rst │ │ ├── frontend.rst │ │ ├── graph.rst │ │ ├── index.rst │ │ ├── symbol.rst │ │ └── top.rst ├── conf.py ├── dev │ ├── index.rst │ └── overview.md ├── how_to │ ├── contribute.md │ ├── deploy.md │ └── install.md ├── index.rst ├── json_spec.rst └── top.rst ├── examples ├── README.md ├── benchmark │ ├── gpu_imagenet_bench.py │ └── rasp_imagenet_bench.py └── sgx │ ├── .gitignore │ ├── Makefile │ ├── README.md │ ├── app.cc │ ├── build_model.py │ ├── enclave.cc │ ├── enclave_config.xml │ ├── enclave_private.pem │ ├── model.edl │ └── run_example.sh ├── include └── nnvm │ ├── base.h │ ├── c_api.h │ ├── compiler │ ├── op_attr_types.h │ ├── packed_func_ext.h │ └── util.h │ ├── graph.h │ ├── graph_attr_types.h │ ├── layout.h │ ├── node.h │ ├── op.h │ ├── op_attr_types.h │ ├── pass.h │ ├── pass_functions.h │ ├── symbolic.h │ ├── top │ ├── README │ ├── nn.h │ └── tensor.h │ └── tuple.h ├── make └── config.mk ├── python ├── .gitignore ├── conda │ ├── build.sh │ └── meta.yaml ├── nnvm │ ├── __init__.py │ ├── _base.py │ ├── _ctypes │ │ ├── README │ │ ├── __init__.py │ │ └── symbol.py │ ├── _cy2 │ │ ├── README │ │ └── __init__.py │ ├── _cy3 │ │ ├── README │ │ └── __init__.py │ ├── _symbol_internal.py │ ├── attribute.py │ ├── compiler │ │ ├── __init__.py │ │ ├── build_module.py │ │ ├── compile_engine.py │ │ ├── graph_attr.py │ │ ├── graph_pass.py │ │ ├── graph_util.py │ │ ├── lr_scheduler.py │ │ ├── optimizer.py │ │ └── param_dict.py │ ├── contrib.py │ ├── cython │ │ ├── README │ │ ├── base.pyi │ │ └── symbol.pyx │ ├── frontend │ │ ├── __init__.py │ │ ├── common.py │ │ ├── coreml.py │ │ ├── darknet.py │ │ ├── keras.py │ │ ├── mxnet.py │ │ └── onnx.py │ ├── graph.py │ ├── libinfo.py │ ├── name.py │ ├── symbol.py │ ├── testing │ │ ├── __init__.py │ │ ├── config.py │ │ ├── darknet.py │ │ ├── init.py │ │ ├── mlp.py │ │ ├── mobilenet.py │ │ ├── resnet.py │ │ ├── utils.py │ │ ├── vgg.py │ │ └── yolo2_detection.py │ └── top │ │ ├── __init__.py │ │ ├── attr_dict.py │ │ ├── nn.py │ │ ├── reduction.py │ │ ├── registry.py │ │ ├── tensor.py │ │ ├── transform.py │ │ └── vision.py └── setup.py ├── src ├── README.md ├── c_api │ ├── c_api_common.h │ ├── c_api_error.cc │ ├── c_api_graph.cc │ └── c_api_symbolic.cc ├── compiler │ ├── alter_op_layout.cc │ ├── compile_engine.cc │ ├── compile_engine.h │ ├── fold_scale_axis.cc │ ├── graph_fuse.cc │ ├── graph_hash.cc │ ├── graph_hash.h │ ├── graph_runtime.cc │ ├── graph_runtime.h │ ├── graph_transform.h │ ├── node_attr.h │ ├── packed_func_ext.cc │ ├── pattern_util.h │ ├── precompute_prune.cc │ └── simplify_inference.cc ├── core │ ├── graph.cc │ ├── node.cc │ ├── op.cc │ ├── pass.cc │ └── symbolic.cc ├── pass │ ├── correct_layout.cc │ ├── gradient.cc │ ├── graph_algorithm.h │ ├── infer_shape_type.cc │ ├── order_mutation.cc │ ├── place_device.cc │ ├── plan_memory.cc │ ├── print_graph_ir.cc │ └── saveload_json.cc └── top │ ├── elemwise_op_common.h │ ├── nn │ ├── convolution.cc │ ├── nn.cc │ ├── nn_common.h │ ├── pooling.cc │ └── upsampling.cc │ ├── op_common.h │ ├── tensor │ ├── broadcast.cc │ ├── elemwise.cc │ ├── matrix_op.cc │ ├── reduce.cc │ ├── state_op.cc │ └── transform.cc │ └── vision │ └── yolo2 │ ├── region.cc │ ├── region.h │ ├── reorg.cc │ └── reorg.h ├── tests ├── ci_build │ ├── Dockerfile.gpu │ ├── Dockerfile.lint │ ├── README.md │ ├── ci_build.sh │ ├── install │ │ ├── ubuntu_install_core.sh │ │ ├── ubuntu_install_coreml.sh │ │ ├── ubuntu_install_darknet.sh │ │ ├── ubuntu_install_keras.sh │ │ ├── ubuntu_install_llvm.sh │ │ ├── ubuntu_install_mxnet.sh │ │ ├── ubuntu_install_onnx.sh │ │ ├── ubuntu_install_opencl.sh │ │ ├── ubuntu_install_python.sh │ │ ├── ubuntu_install_python_package.sh │ │ └── ubuntu_install_sphinx.sh │ └── with_the_same_user ├── cpp │ ├── .gitignore │ ├── op_test.cc │ ├── tuple_test.cc │ └── unittest.mk ├── lint │ └── pylintrc ├── python │ ├── compiler │ │ ├── test_alter_op_layout.py │ │ ├── test_build.py │ │ ├── test_compiler_cache.py │ │ ├── test_fold_axis.py │ │ ├── test_graph_pass.py │ │ ├── test_nhwc_layout.py │ │ ├── test_op_fusion.py │ │ ├── test_optimizer.py │ │ ├── test_param_dict.py │ │ ├── test_rpc_exec.py │ │ ├── test_simplify_inference.py │ │ ├── test_top_assign.py │ │ ├── test_top_level1.py │ │ ├── test_top_level2.py │ │ └── test_top_level4.py │ ├── frontend │ │ ├── coreml │ │ │ ├── model_zoo │ │ │ │ ├── .gitignore │ │ │ │ └── __init__.py │ │ │ └── test_forward.py │ │ ├── darknet │ │ │ └── test_forward.py │ │ ├── keras │ │ │ └── test_forward.py │ │ ├── mxnet │ │ │ ├── model_zoo │ │ │ │ ├── __init__.py │ │ │ │ ├── mlp.py │ │ │ │ ├── resnet.py │ │ │ │ └── vgg.py │ │ │ ├── test_forward.py │ │ │ └── test_graph.py │ │ └── onnx │ │ │ ├── model_zoo │ │ │ ├── __init__.py │ │ │ └── super_resolution.py │ │ │ ├── test_forward.py │ │ │ └── test_graph.py │ └── unittest │ │ ├── test_correct_layout.py │ │ ├── test_graph.py │ │ ├── test_graph_gradient.py │ │ ├── test_infer_shape.py │ │ ├── test_symbol.py │ │ ├── test_top_level1.py │ │ ├── test_top_level2.py │ │ ├── test_top_level3.py │ │ └── test_top_level4.py ├── scripts │ ├── task_build.sh │ ├── task_clean.sh │ ├── task_frontend_test.sh │ ├── task_lint.sh │ ├── task_python_docs.sh │ └── task_python_test.sh └── travis │ ├── run_test.sh │ ├── setup.sh │ └── travis_after_failure.sh └── tutorials ├── README.txt ├── define_and_compile_model.py ├── deploy_model_on_mali_gpu.py ├── deploy_model_on_rasp.py ├── from_coreml.py ├── from_darknet.py ├── from_keras.py ├── from_mxnet.py ├── from_mxnet_to_webgl.py ├── from_onnx.py ├── get_started.py ├── imagenet_inference_gpu.py ├── using_external_lib.py └── web └── resnet.html /.gitignore: -------------------------------------------------------------------------------- 1 | # dmlc 2 | /config.mk 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | build 34 | lib 35 | *~ 36 | cli_test 37 | *.pyc 38 | .nv 39 | .cache 40 | docs.tgz 41 | # Vim 42 | *.swp 43 | *.swo 44 | 45 | # TVM generated code 46 | perf 47 | 48 | *.json 49 | *.params 50 | *.onnx 51 | *.h5 52 | synset.txt 53 | cat.jpg 54 | 55 | # Mac OS X 56 | .DS_Store 57 | 58 | # Jetbrain 59 | .idea 60 | 61 | # Distribution / packaging 62 | *.egg-info/ 63 | dist/ 64 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "plugin/nnvm-fusion"] 2 | path = plugin/nnvm-fusion 3 | url = https://github.com/dmlc/nnvm-fusion.git 4 | [submodule "dmlc-core"] 5 | path = dmlc-core 6 | url = https://github.com/dmlc/dmlc-core 7 | [submodule "tvm"] 8 | path = tvm 9 | url = https://github.com/dmlc/tvm 10 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | sudo: false 2 | 3 | language: cpp 4 | 5 | os: 6 | - linux 7 | - osx 8 | 9 | osx_image: xcode8 10 | 11 | env: 12 | # code analysis 13 | - TASK=lint 14 | - TASK=cpp_test 15 | - TASK=python_test 16 | 17 | matrix: 18 | exclude: 19 | - os: osx 20 | env: TASK=lint 21 | 22 | # dependent apt packages 23 | addons: 24 | apt: 25 | sources: 26 | - ubuntu-toolchain-r-test 27 | packages: 28 | - doxygen 29 | - wget 30 | - git 31 | - unzip 32 | - gcc-4.8 33 | - g++-4.8 34 | - python-numpy 35 | - python-nose 36 | - python3-numpy 37 | - python3-dev 38 | - python3-nose 39 | - graphviz 40 | 41 | before_install: 42 | - source dmlc-core/scripts/travis/travis_setup_env.sh 43 | - export PYTHONPATH=${PYTHONPATH}:${PWD}/python 44 | 45 | install: 46 | - source tests/travis/setup.sh 47 | 48 | script: 49 | - tests/travis/run_test.sh 50 | 51 | cache: 52 | directories: 53 | - ${HOME}/.cache/usr 54 | 55 | before_cache: 56 | - dmlc-core/scripts/travis/travis_before_cache.sh 57 | 58 | after_failure: 59 | - tests/travis/travis_after_failure.sh 60 | 61 | notifications: 62 | # Emails are sent to the committer's git-configured email address by default, 63 | email: 64 | on_success: change 65 | on_failure: always 66 | #slack: dmlc:NmroCzntCiWOuxUZpii40USd 67 | -------------------------------------------------------------------------------- /Jenkinsfile: -------------------------------------------------------------------------------- 1 | #!groovy 2 | // -*- mode: groovy -*- 3 | // Jenkins pipeline 4 | // See documents at https://jenkins.io/doc/book/pipeline/jenkinsfile/ 5 | 6 | // nnvm libraries 7 | nnvm_lib = "tvm/lib/libtvm.so, tvm/lib/libtvm_runtime.so, lib/libnnvm_compiler.so" 8 | 9 | // command to start a docker container 10 | docker_run = 'tests/ci_build/ci_build.sh' 11 | // timeout in minutes 12 | max_time = 60 13 | 14 | // initialize source codes 15 | def init_git() { 16 | checkout scm 17 | retry(5) { 18 | timeout(time: 2, unit: 'MINUTES') { 19 | sh 'git submodule update --init --recursive' 20 | } 21 | } 22 | } 23 | 24 | def init_git_win() { 25 | checkout scm 26 | retry(5) { 27 | timeout(time: 2, unit: 'MINUTES') { 28 | bat 'git submodule update --init --recursive' 29 | } 30 | } 31 | } 32 | 33 | stage("Sanity Check") { 34 | timeout(time: max_time, unit: 'MINUTES') { 35 | node('linux') { 36 | ws('workspace/tvm/sanity') { 37 | init_git() 38 | sh "${docker_run} lint ./tests/scripts/task_lint.sh" 39 | } 40 | } 41 | } 42 | } 43 | 44 | // Run make. First try to do an incremental make from a previous workspace in hope to 45 | // accelerate the compilation. If something wrong, clean the workspace and then 46 | // build from scratch. 47 | def make(docker_type, make_flag) { 48 | timeout(time: max_time, unit: 'MINUTES') { 49 | try { 50 | sh "${docker_run} ${docker_type} ./tests/scripts/task_build.sh ${make_flag}" 51 | } catch (exc) { 52 | echo 'Incremental compilation failed. Fall back to build from scratch' 53 | sh "${docker_run} ${docker_type} ./tests/scripts/task_clean.sh" 54 | sh "${docker_run} ${docker_type} ./tests/scripts/task_build.sh ${make_flag}" 55 | } 56 | } 57 | } 58 | 59 | // pack libraries for later use 60 | def pack_lib(name, libs) { 61 | sh """ 62 | echo "Packing ${libs} into ${name}" 63 | echo ${libs} | sed -e 's/,/ /g' | xargs md5sum 64 | """ 65 | stash includes: libs, name: name 66 | } 67 | 68 | 69 | // unpack libraries saved before 70 | def unpack_lib(name, libs) { 71 | unstash name 72 | sh """ 73 | echo "Unpacked ${libs} from ${name}" 74 | echo ${libs} | sed -e 's/,/ /g' | xargs md5sum 75 | """ 76 | } 77 | 78 | stage('Build') { 79 | timeout(time: max_time, unit: 'MINUTES') { 80 | node('GPU' && 'linux') { 81 | ws('workspace/nnvm/build-gpu') { 82 | init_git() 83 | make('gpu', '-j2') 84 | pack_lib('gpu', nnvm_lib) 85 | } 86 | } 87 | } 88 | } 89 | 90 | stage('Tests') { 91 | parallel 'python': { 92 | node('GPU' && 'linux') { 93 | ws('workspace/nnvm/it-python-gpu') { 94 | init_git() 95 | unpack_lib('gpu', nnvm_lib) 96 | timeout(time: max_time, unit: 'MINUTES') { 97 | sh "${docker_run} gpu ./tests/scripts/task_python_test.sh" 98 | sh "${docker_run} gpu ./tests/scripts/task_frontend_test.sh" 99 | } 100 | } 101 | } 102 | }, 103 | 'docs': { 104 | node('GPU' && 'linux') { 105 | ws('workspace/nnvm/docs-python-gpu') { 106 | init_git() 107 | unpack_lib('gpu', nnvm_lib) 108 | timeout(time: max_time, unit: 'MINUTES') { 109 | sh "${docker_run} gpu ./tests/scripts/task_python_docs.sh" 110 | } 111 | pack_lib('mydocs', 'docs.tgz') 112 | } 113 | } 114 | } 115 | } 116 | 117 | stage('Deploy') { 118 | node('docker' && 'doc') { 119 | ws('workspace/nnvm/deploy-docs') { 120 | if (env.BRANCH_NAME == "master") { 121 | unpack_lib('mydocs', 'docs.tgz') 122 | sh "tar xf docs.tgz -C /var/nnvm-docs" 123 | } 124 | } 125 | } 126 | } 127 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | ROOTDIR = $(CURDIR) 2 | 3 | ifndef config 4 | ifneq ("$(wildcard ./config.mk)", "") 5 | config = config.mk 6 | else 7 | config = make/config.mk 8 | endif 9 | endif 10 | include $(config) 11 | 12 | export LDFLAGS = -pthread -lm 13 | export CFLAGS = -std=c++11 -Wall -O2 -Iinclude -fPIC 14 | CFLAGS += -Itvm/include -Itvm/dlpack/include -Itvm/HalideIR/src -Itvm/topi/include 15 | 16 | ifdef DMLC_CORE_PATH 17 | CFLAGS += -I$(DMLC_CORE_PATH)/include 18 | else 19 | CFLAGS += -I$(ROOTDIR)/dmlc-core/include 20 | endif 21 | 22 | ifdef DMLC_CORE_PATH 23 | CFLAGS += -I$(DMLC_CORE_PATH)/include 24 | else 25 | CFLAGS += -I$(ROOTDIR)/dmlc-core/include 26 | endif 27 | 28 | ifneq ($(ADD_CFLAGS), NONE) 29 | CFLAGS += $(ADD_CFLAGS) 30 | endif 31 | 32 | ifneq ($(ADD_LDFLAGS), NONE) 33 | LDFLAGS += $(ADD_LDFLAGS) 34 | endif 35 | 36 | # plugin 37 | PLUGIN_OBJ = 38 | include $(NNVM_PLUGINS) 39 | 40 | # specify tensor path 41 | .PHONY: clean all test lint cpplint pylint doc cython cython3 cyclean 42 | 43 | UNAME_S := $(shell uname -s) 44 | 45 | ifeq ($(UNAME_S), Darwin) 46 | SHARED_LIBRARY_SUFFIX := dylib 47 | WHOLE_ARCH= -all_load 48 | NO_WHOLE_ARCH= -noall_load 49 | LDFLAGS += -undefined dynamic_lookup 50 | else 51 | SHARED_LIBRARY_SUFFIX := so 52 | WHOLE_ARCH= --whole-archive 53 | NO_WHOLE_ARCH= --no-whole-archive 54 | endif 55 | 56 | all: lib/libnnvm.a lib/libnnvm_compiler.$(SHARED_LIBRARY_SUFFIX) 57 | 58 | SRC = $(wildcard src/*.cc src/c_api/*.cc src/core/*.cc src/pass/*.cc) 59 | SRC_COMPILER = $(wildcard src/top/*/*.cc wildcard src/top/vision/*/*.cc src/compiler/*.cc src/compiler/*/*.cc) 60 | ALL_OBJ = $(patsubst %.cc, build/%.o, $(SRC)) 61 | TOP_OBJ = $(patsubst %.cc, build/%.o, $(SRC_COMPILER)) 62 | ALL_DEP = $(ALL_OBJ) 63 | 64 | include tests/cpp/unittest.mk 65 | 66 | test: $(TEST) 67 | 68 | build/src/%.o: src/%.cc 69 | @mkdir -p $(@D) 70 | $(CXX) $(CFLAGS) -MM -MT build/src/$*.o $< >build/src/$*.d 71 | $(CXX) -c $(CFLAGS) -c $< -o $@ 72 | 73 | lib/libnnvm.a: $(ALL_DEP) 74 | @mkdir -p $(@D) 75 | $(AR) crv $@ $(filter %.o, $?) 76 | 77 | lib/libnnvm_compiler.$(SHARED_LIBRARY_SUFFIX): lib/libnnvm.a ${TOP_OBJ} 78 | @mkdir -p $(@D) 79 | $(CXX) $(CFLAGS) -shared -o $@ $(filter %.o, $^) $(LDFLAGS) -Wl,${WHOLE_ARCH} lib/libnnvm.a -Wl,${NO_WHOLE_ARCH} 80 | 81 | cython: 82 | cd python; python setup.py build_ext --inplace 83 | 84 | cython3: 85 | cd python; python3 setup.py build_ext --inplace 86 | 87 | cyclean: 88 | rm -rf python/nnvm/*/*.so python/nnvm/*/*.dylib python/nnvm/*/*.cpp 89 | 90 | lint: pylint cpplint 91 | 92 | cpplint: 93 | python dmlc-core/scripts/lint.py nnvm cpp include src 94 | 95 | pylint: 96 | pylint python/nnvm --rcfile=$(ROOTDIR)/tests/lint/pylintrc 97 | 98 | doc: 99 | doxygen docs/Doxyfile 100 | 101 | clean: 102 | $(RM) -rf build lib bin *~ */*~ */*/*~ */*/*/*~ */*.o */*/*.o */*/*/*.o cli_test 103 | 104 | -include build/*.d 105 | -include build/*/*.d 106 | -include build/*/*/*.d 107 | -include build/*/*/*/*.d 108 | -------------------------------------------------------------------------------- /NEWS.md: -------------------------------------------------------------------------------- 1 | NNVM Change Log 2 | =============== 3 | 4 | This file records the changes in TVM library in reverse chronological order. 5 | 6 | ## 0.8rc 7 | 8 | - This is major change in NNVM to introduce end to end compiler stack. 9 | - The NNVM compiler stack ready 10 | - Core tensor operators 11 | - integrates compiler with TVM 12 | - The libnnvm.a is still independent from compiler modules. 13 | 14 | ## 0.7 15 | 16 | - NNVM graph 17 | - Basic pass of serialization, gradient, infer_shape, place_deice, plan_memory 18 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # [Source code now lives in the TVM repository](https://github.com/dmlc/tvm/). 2 | 3 | -------------------------------------------------------------------------------- /amalgamation/.gitignore: -------------------------------------------------------------------------------- 1 | nnvm.d 2 | nnvm.cc 3 | -------------------------------------------------------------------------------- /amalgamation/Makefile: -------------------------------------------------------------------------------- 1 | export NNVM_ROOT=`pwd`/.. 2 | export CFLAGS = -std=c++11 -Wall -O2 -Iinclude -fPIC 3 | 4 | ifdef DMLC_CORE_PATH 5 | CFLAGS += -I$(DMLC_CORE_PATH)/include 6 | else 7 | CFLAGS += -I$(CURDIR)/../dmlc-core/include 8 | endif 9 | 10 | .PHONY: all clean 11 | 12 | all: libnnvm.a 13 | 14 | nnvm.cc: 15 | python generate.py $@ 16 | 17 | nnvm.d: nnvm.cc 18 | ${CXX} ${CFLAGS} -M -MT nnvm.o \ 19 | -I ${NNVM_ROOT}/ -I ${NNVM_ROOT}/include \ 20 | -D__MIN__=$(MIN) $+ > nnvm.d 21 | 22 | nnvm-all.cc: nnvm.d nnvm.cc 23 | python ./amalgamation.py $+ $@ 24 | 25 | nnvm-all.o: nnvm-all.cc 26 | ${CXX} ${CFLAGS} -fPIC -o $@ -c $+ 27 | 28 | libnnvm.a: nnvm-all.o 29 | ar rcs $@ $+ 30 | 31 | clean: 32 | rm -f *.d *.o *.so *.a nnvm-all.cc nnvm.cc 33 | -------------------------------------------------------------------------------- /amalgamation/README: -------------------------------------------------------------------------------- 1 | This folder is deprecated and will be deleted in the future. -------------------------------------------------------------------------------- /amalgamation/amalgamation.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os.path, re, StringIO 3 | 4 | blacklist = [ 5 | 'Windows.h', 6 | 'mach/clock.h', 'mach/mach.h', 7 | 'malloc.h', 8 | 'glog/logging.h', 'io/azure_filesys.h', 'io/hdfs_filesys.h', 'io/s3_filesys.h', 9 | 'sys/stat.h', 'sys/types.h', 10 | 'omp.h', 'execinfo.h', 'packet/sse-inl.h' 11 | ] 12 | 13 | 14 | def get_sources(def_file): 15 | sources = [] 16 | files = [] 17 | visited = set() 18 | mxnet_path = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)) 19 | for line in open(def_file): 20 | files = files + line.strip().split(' ') 21 | 22 | for f in files: 23 | f = f.strip() 24 | if not f or f.endswith('.o:') or f == '\\': continue 25 | fn = os.path.relpath(f) 26 | if os.path.abspath(f).startswith(mxnet_path) and fn not in visited: 27 | sources.append(fn) 28 | visited.add(fn) 29 | return sources 30 | 31 | sources = get_sources(sys.argv[1]) 32 | 33 | def find_source(name, start): 34 | candidates = [] 35 | for x in sources: 36 | if x == name or x.endswith('/' + name): candidates.append(x) 37 | if not candidates: return '' 38 | if len(candidates) == 1: return candidates[0] 39 | for x in candidates: 40 | if x.split('/')[1] == start.split('/')[1]: return x 41 | return '' 42 | 43 | 44 | re1 = re.compile('<([./a-zA-Z0-9_-]*)>') 45 | re2 = re.compile('"([./a-zA-Z0-9_-]*)"') 46 | 47 | sysheaders = [] 48 | history = set([]) 49 | out = StringIO.StringIO() 50 | 51 | def expand(x, pending): 52 | if x in history and x not in ['mshadow/mshadow/expr_scalar-inl.h']: # MULTIPLE includes 53 | return 54 | 55 | if x in pending: 56 | #print 'loop found: %s in ' % x, pending 57 | return 58 | 59 | print >>out, "//===== EXPANDING: %s =====\n" %x 60 | for line in open(x): 61 | if line.find('#include') < 0: 62 | out.write(line) 63 | continue 64 | if line.strip().find('#include') > 0: 65 | print line 66 | continue 67 | m = re1.search(line) 68 | if not m: m = re2.search(line) 69 | if not m: 70 | print line + ' not found' 71 | continue 72 | h = m.groups()[0].strip('./') 73 | source = find_source(h, x) 74 | if not source: 75 | if (h not in blacklist and 76 | h not in sysheaders and 77 | 'mkl' not in h and 78 | 'nnpack' not in h): sysheaders.append(h) 79 | else: 80 | expand(source, pending + [x]) 81 | print >>out, "//===== EXPANDED: %s =====\n" %x 82 | history.add(x) 83 | 84 | 85 | expand(sys.argv[2], []) 86 | 87 | f = open(sys.argv[3], 'wb') 88 | 89 | 90 | 91 | for k in sorted(sysheaders): 92 | print >>f, "#include <%s>" % k 93 | 94 | print >>f, '' 95 | print >>f, out.getvalue() 96 | 97 | for x in sources: 98 | if x not in history and not x.endswith('.o'): 99 | print 'Not processed:', x 100 | 101 | -------------------------------------------------------------------------------- /amalgamation/generate.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | FOLDERS = ["core", "pass", "c_api"] 5 | 6 | fo = open(sys.argv[1], "w") 7 | 8 | 9 | 10 | for folder in FOLDERS: 11 | path = str(os.path.join("../src", folder)) 12 | flst = os.listdir(path) 13 | for f in flst: 14 | if f.endswith(".cc") == True: 15 | fo.write('#include "' + str(os.path.join("src", folder, f)) + '"\n') 16 | 17 | 18 | fo.close() 19 | -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | doxygen 2 | _build 3 | gen_modules 4 | tutorials 5 | -------------------------------------------------------------------------------- /docs/README.txt: -------------------------------------------------------------------------------- 1 | The documentation of nnvm is generated with recommonmark and sphinx. 2 | 3 | - pip install sphinx>=1.5.5 sphinx-gallery sphinx_rtd_theme matplotlib Image recommonmark 4 | - Build tvm first in the root folder. 5 | -------------------------------------------------------------------------------- /docs/api/python/compiler.rst: -------------------------------------------------------------------------------- 1 | nnvm.compiler 2 | ------------- 3 | 4 | .. automodule:: nnvm.compiler 5 | 6 | .. autofunction:: nnvm.compiler.build 7 | 8 | .. autofunction:: nnvm.compiler.build_config 9 | 10 | .. autofunction:: nnvm.compiler.save_param_dict 11 | 12 | .. autofunction:: nnvm.compiler.load_param_dict 13 | 14 | .. autofunction:: nnvm.compiler.optimize 15 | 16 | .. automodule:: nnvm.compiler.graph_util 17 | :members: 18 | 19 | .. automodule:: nnvm.compiler.graph_attr 20 | :members: 21 | 22 | .. automodule:: nnvm.compiler.compile_engine 23 | :members: 24 | -------------------------------------------------------------------------------- /docs/api/python/frontend.rst: -------------------------------------------------------------------------------- 1 | nnvm.frontend 2 | ------------- 3 | 4 | .. automodule:: nnvm.frontend 5 | 6 | .. autofunction:: nnvm.frontend.from_mxnet 7 | 8 | .. autofunction:: nnvm.frontend.from_onnx 9 | 10 | .. autofunction:: nnvm.frontend.from_coreml 11 | 12 | .. autofunction:: nnvm.frontend.from_keras 13 | -------------------------------------------------------------------------------- /docs/api/python/graph.rst: -------------------------------------------------------------------------------- 1 | nnvm.graph 2 | ---------- 3 | .. automodule:: nnvm.graph 4 | 5 | .. autofunction:: nnvm.graph.create 6 | 7 | .. autoclass:: nnvm.graph.Graph 8 | :members: 9 | -------------------------------------------------------------------------------- /docs/api/python/index.rst: -------------------------------------------------------------------------------- 1 | Python API 2 | ========== 3 | 4 | This document contains the python API to NNVM compiler toolchain. 5 | For user 6 | 7 | 8 | .. toctree:: 9 | :maxdepth: 2 10 | 11 | compiler 12 | frontend 13 | symbol 14 | graph 15 | top 16 | -------------------------------------------------------------------------------- /docs/api/python/symbol.rst: -------------------------------------------------------------------------------- 1 | nnvm.symbol 2 | ----------- 3 | .. automodule:: nnvm.symbol 4 | 5 | .. autoclass:: nnvm.symbol.Symbol 6 | :members: 7 | 8 | .. autoclass:: nnvm.symbol.Variable 9 | 10 | .. autofunction:: nnvm.symbol.Group 11 | -------------------------------------------------------------------------------- /docs/api/python/top.rst: -------------------------------------------------------------------------------- 1 | nnvm.top 2 | -------- 3 | .. automodule:: nnvm.top 4 | 5 | .. autofunction:: register_compute 6 | 7 | .. autofunction:: register_schedule 8 | 9 | .. autofunction:: register_pattern 10 | 11 | 12 | .. autoclass:: nnvm.top.AttrDict 13 | :members: 14 | -------------------------------------------------------------------------------- /docs/dev/index.rst: -------------------------------------------------------------------------------- 1 | Design Note 2 | =========== 3 | 4 | In this part of documentation, we share the rationale for the specific choices made when designing NNVM. 5 | 6 | .. toctree:: 7 | :maxdepth: 2 8 | 9 | overview 10 | -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | NNVM Documentation 2 | ================== 3 | This is a document about NNVM and NNVM compiler. 4 | 5 | Contents 6 | -------- 7 | 8 | .. toctree:: 9 | :maxdepth: 1 10 | 11 | self 12 | how_to/install 13 | tutorials/index 14 | top 15 | json_spec 16 | how_to/contribute 17 | how_to/deploy 18 | api/python/index 19 | dev/index 20 | -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- 1 | NNVM Examples 2 | ============= 3 | This folder contains example snippets of running NNVM Compilation. 4 | 5 | - See also [Tutorials](../tutorials) for tutorials with detailed explainations. 6 | -------------------------------------------------------------------------------- /examples/benchmark/gpu_imagenet_bench.py: -------------------------------------------------------------------------------- 1 | """ Benchmark script for performance on GPUs. 2 | 3 | For example, run the file with: 4 | `python gpu_imagenet_bench.py --model=mobilenet --target=cuda`. 5 | For more details about how to set up the inference environment on GPUs, 6 | please refer to NNVM Tutorial: ImageNet Inference on the GPU 7 | """ 8 | import time 9 | import argparse 10 | import numpy as np 11 | import tvm 12 | import nnvm.compiler 13 | import nnvm.testing 14 | from tvm.contrib import util, nvcc 15 | from tvm.contrib import graph_runtime as runtime 16 | 17 | @tvm.register_func 18 | def tvm_callback_cuda_compile(code): 19 | ptx = nvcc.compile_cuda(code, target="ptx") 20 | return ptx 21 | 22 | def main(): 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--model', type=str, required=True, 25 | choices=['resnet', 'mobilenet'], 26 | help="The model type.") 27 | parser.add_argument('--target', type=str, required=True, 28 | choices=['cuda', 'rocm', 'opencl', 'metal'], 29 | help="Compilation target.") 30 | parser.add_argument('--opt-level', type=int, default=1, help="Level of optimization.") 31 | parser.add_argument('--num-iter', type=int, default=1000, help="Number of iteration during benchmark.") 32 | parser.add_argument('--repeat', type=int, default=1, help="Number of repeative times.") 33 | args = parser.parse_args() 34 | opt_level = args.opt_level 35 | num_iter = args.num_iter 36 | ctx = tvm.context(args.target, 0) 37 | batch_size = 1 38 | num_classes = 1000 39 | image_shape = (3, 224, 224) 40 | 41 | data_shape = (batch_size,) + image_shape 42 | out_shape = (batch_size, num_classes) 43 | if args.model == 'resnet': 44 | net, params = nnvm.testing.resnet.get_workload( 45 | batch_size=1, image_shape=image_shape) 46 | elif args.model == 'mobilenet': 47 | net, params = nnvm.testing.mobilenet.get_workload( 48 | batch_size=1, image_shape=image_shape) 49 | else: 50 | raise ValueError('no benchmark prepared for {}.'.format(args.model)) 51 | 52 | if args.target == "cuda": 53 | unroll = 1400 54 | else: 55 | unroll = 128 56 | with nnvm.compiler.build_config(opt_level=opt_level): 57 | with tvm.build_config(auto_unroll_max_step=unroll, 58 | unroll_explicit=(args.target != "cuda")): 59 | graph, lib, params = nnvm.compiler.build( 60 | net, args.target, shape={"data": data_shape}, params=params) 61 | 62 | data = np.random.uniform(-1, 1, size=data_shape).astype("float32") 63 | module = runtime.create(graph, lib, ctx) 64 | module.set_input(**params) 65 | module.set_input("data", data) 66 | module.run() 67 | out = module.get_output(0, tvm.nd.empty(out_shape)) 68 | out.asnumpy() 69 | 70 | print('benchmark args: {}'.format(args)) 71 | ftimer = module.module.time_evaluator("run", ctx, num_iter) 72 | for i in range(args.repeat): 73 | prof_res = ftimer() 74 | print(prof_res) 75 | # sleep for avoiding device overheat 76 | if i + 1 != args.repeat: 77 | time.sleep(45) 78 | 79 | if __name__ == '__main__': 80 | main() 81 | -------------------------------------------------------------------------------- /examples/benchmark/rasp_imagenet_bench.py: -------------------------------------------------------------------------------- 1 | """ Benchmark script for performance on Raspberry Pi. For example, run the file with: 2 | `python rasp_imagenet_bench.py --model='modbilenet' --host='rasp0' --port=9090`. For 3 | more details about how to set up the inference environment on Raspberry Pi, Please 4 | refer to NNVM Tutorial: Deploy the Pretrained Model on Raspberry Pi """ 5 | import time 6 | import argparse 7 | import numpy as np 8 | import tvm 9 | import nnvm.compiler 10 | import nnvm.testing 11 | from tvm.contrib import util, rpc 12 | from tvm.contrib import graph_runtime as runtime 13 | 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--model', type=str, required=True, choices=['resnet', 'mobilenet'], 18 | help="The model type.") 19 | parser.add_argument('--host', type=str, required=True, help="The host address of your Raspberry Pi.") 20 | parser.add_argument('--port', type=int, required=True, help="The port number of your Raspberry Pi.") 21 | parser.add_argument('--opt-level', type=int, default=1, help="Level of optimization.") 22 | parser.add_argument('--num-iter', type=int, default=50, help="Number of iteration during benchmark.") 23 | args = parser.parse_args() 24 | 25 | opt_level = args.opt_level 26 | 27 | num_iter = args.num_iter 28 | batch_size = 1 29 | num_classes = 1000 30 | image_shape = (3, 224, 224) 31 | 32 | data_shape = (batch_size,) + image_shape 33 | out_shape = (batch_size, num_classes) 34 | if args.model == 'resnet': 35 | net, params = nnvm.testing.resnet.get_workload( 36 | batch_size=1, image_shape=image_shape) 37 | elif args.model == 'mobilenet': 38 | net, params = nnvm.testing.mobilenet.get_workload( 39 | batch_size=1, image_shape=image_shape) 40 | else: 41 | raise ValueError('no benchmark prepared for {}.'.format(args.model)) 42 | 43 | 44 | with nnvm.compiler.build_config(opt_level=opt_level): 45 | graph, lib, params = nnvm.compiler.build( 46 | net, tvm.target.rasp(), shape={"data": data_shape}, params=params) 47 | 48 | tmp = util.tempdir() 49 | lib_fname = tmp.relpath('net.o') 50 | lib.save(lib_fname) 51 | 52 | remote = rpc.connect(args.host, args.port) 53 | remote.upload(lib_fname) 54 | 55 | ctx = remote.cpu(0) 56 | rlib = remote.load_module('net.o') 57 | rparams = {k: tvm.nd.array(v, ctx) for k, v in params.items()} 58 | 59 | module = runtime.create(graph, rlib, ctx) 60 | module.set_input('data', tvm.nd.array(np.random.uniform(size=(data_shape)).astype("float32"))) 61 | module.set_input(**rparams) 62 | module.run() 63 | out = module.get_output(0, tvm.nd.empty(out_shape, ctx=ctx)) 64 | out.asnumpy() 65 | 66 | print('benchmark args: {}'.format(args)) 67 | ftimer = module.module.time_evaluator("run", ctx, num_iter) 68 | for i in range(3): 69 | prof_res = ftimer() 70 | print(prof_res) 71 | # sleep for avoiding cpu overheat 72 | time.sleep(45) 73 | 74 | 75 | if __name__ == '__main__': 76 | main() 77 | -------------------------------------------------------------------------------- /examples/sgx/.gitignore: -------------------------------------------------------------------------------- 1 | lib/ 2 | bin/ 3 | -------------------------------------------------------------------------------- /examples/sgx/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for example to deploy TVM modules in SGX. 2 | 3 | PYTHON ?= python 4 | 5 | NNVM_ROOT := $(shell cd ../../; pwd) 6 | TVM_ROOT := $(NNVM_ROOT)/tvm 7 | DMLC_CORE_ROOT := $(NNVM_ROOT)/dmlc-core 8 | 9 | SGX_SDK ?= /opt/sgxsdk 10 | SGX_MODE ?= SIM 11 | SGX_ARCH ?= x64 12 | SGX_DEBUG ?= 1 13 | 14 | sgx_edger8r := $(SGX_SDK)/bin/x64/sgx_edger8r 15 | sgx_enclave_signer := $(SGX_SDK)/bin/x64/sgx_sign 16 | 17 | ifneq ($(SGX_MODE), HW) 18 | sgx_sim := _sim 19 | endif 20 | urts_library_name := sgx_urts$(sgx_sim) 21 | trts_library_name := sgx_trts$(sgx_sim) 22 | tservice_library_name := sgx_tservice$(sgx_sim) 23 | uservice_library_name := sgx_uae_service$(sgx_sim) 24 | 25 | pkg_cflags := -std=c++11 -O2 -fPIC\ 26 | -I$(NNVM_ROOT)/include\ 27 | -I$(NNVM_ROOT)\ 28 | -I$(TVM_ROOT)/include\ 29 | -I$(TVM_ROOT)/dlpack/include\ 30 | -I$(DMLC_CORE_ROOT)/include\ 31 | -DDMLC_LOG_STACK_TRACE=0\ 32 | 33 | pkg_ldflags := -L$(TVM_ROOT)/lib 34 | 35 | enclave_include_paths := -I$(SGX_SDK)/include\ 36 | -I$(SGX_SDK)/include/tlibc\ 37 | -I$(SGX_SDK)/include/libcxx\ 38 | -I$(SGX_SDK)/include/stdc++\ 39 | 40 | enclave_cflags := -static -nostdinc\ 41 | -fvisibility=hidden -fpie -fstack-protector-strong\ 42 | -ffunction-sections -fdata-sections\ 43 | -DDMLC_CXX11_THREAD_LOCAL=0\ 44 | $(enclave_include_paths)\ 45 | 46 | enclave_cxxflags := -nostdinc++ $(enclave_cflags) 47 | 48 | enclave_ldflags :=\ 49 | -Wl,--no-undefined -nostdlib -nodefaultlibs -nostartfiles -L$(SGX_SDK)/lib64\ 50 | -Wl,--whole-archive -l$(trts_library_name) -Wl,--no-whole-archive\ 51 | -Wl,--start-group\ 52 | -lsgx_tstdc -lsgx_tstdcxx -lsgx_tcxx -lsgx_tcrypto -lsgx_tkey_exchange -l$(tservice_library_name)\ 53 | -Wl,--end-group\ 54 | -Wl,-Bstatic -Wl,-Bsymbolic -Wl,--no-undefined\ 55 | -Wl,-pie,-eenclave_entry -Wl,--export-dynamic\ 56 | -Wl,--defsym,__ImageBase=0 -Wl,--gc-sections 57 | 58 | app_cflags := -I$(SGX_SDK)/include -Ilib 59 | 60 | app_ldflags := -L$(SGX_SDK)/lib64\ 61 | -l$(urts_library_name) -l$(uservice_library_name) -lpthread\ 62 | 63 | .PHONY: clean all 64 | 65 | all: lib/model.signed.so bin/run_model 66 | 67 | # The code library built by TVM 68 | lib/deploy_%.o: build_model.py 69 | @mkdir -p $(@D) 70 | $(PYTHON) build_model.py 71 | 72 | # EDL files 73 | lib/model_%.c: model.edl $(sgx_edger8r) 74 | @mkdir -p $(@D) 75 | $(sgx_edger8r) $< --trusted-dir $(@D) --untrusted-dir $(@D) --search-path $(SGX_SDK)/include 76 | 77 | lib/model_%.o: lib/model_%.c 78 | $(CC) $(enclave_cflags) -c $< -o $@ 79 | 80 | # The enclave library 81 | lib/model.so: enclave.cc $(TVM_ROOT)/sgx/sgx_runtime.cc lib/model_t.o lib/deploy_lib.o 82 | $(CXX) $^ -o $@ $(pkg_cflags) $(pkg_ldflags) $(enclave_cxxflags) $(enclave_ldflags)\ 83 | -Wl,--format=binary -Wl,lib/deploy_graph.json -Wl,lib/deploy_params.bin -Wl,--format=default 84 | 85 | # The signed enclave 86 | lib/model.signed.so: lib/model.so enclave_config.xml 87 | $(sgx_enclave_signer) sign -key enclave_private.pem -enclave $< -out $@ -config enclave_config.xml 88 | 89 | # An app that runs the enclave 90 | bin/run_model: app.cc lib/model_u.o 91 | @mkdir -p $(@D) 92 | $(CXX) $^ -o $@ $(app_cflags) $(app_ldflags) 93 | 94 | # Debugging binary that runs TVM without SGX 95 | bin/run_model_nosgx: enclave.cc $(TVM_ROOT)/sgx/sgx_runtime.cc lib/deploy_lib.o 96 | @mkdir -p $(@D) 97 | $(CXX) $^ -o $@ $(pkg_cflags) $(pkg_ldflags)\ 98 | -Wl,--format=binary -Wl,lib/deploy_graph.json -Wl,lib/deploy_params.bin -Wl,--format=default 99 | 100 | 101 | clean: 102 | rm -rf lib bin 103 | -------------------------------------------------------------------------------- /examples/sgx/README.md: -------------------------------------------------------------------------------- 1 | # TVM in Intel SGX Example 2 | 3 | This application demonstrates running a ResNet18 using NNVM inside of an 4 | [Intel SGX](https://software.intel.com/en-us/blogs/2013/09/26/protecting-application-secrets-with-intel-sgx) trusted computing environment. 5 | 6 | ## Prerequisites 7 | 8 | 1. A GNU/Linux environment 9 | 2. NNVM, TVM compiled with LLVM, and their corresponding Python modules 10 | 3. The [Linux SGX SDK](https://github.com/intel/linux-sgx) [link to pre-built libraries](https://01.org/intel-software-guard-extensions/downloads) 11 | 4. `pip install --user mxnet pillow` 12 | 13 | ## Running the example 14 | 15 | `SGX_SDK=/path/to/sgxsdk bash run_example.sh` 16 | 17 | If everything goes well, you should see a lot of build messages and below them 18 | the text `It's a tabby!`. 19 | 20 | ## High-level overview 21 | 22 | First of all, it helps to think of an SGX enclave as a library that can be called 23 | to perform trusted computation. 24 | In this library, one can use other libraries like TVM. 25 | 26 | Building this example performs the following steps: 27 | 28 | 1. Downloads a pre-trained MXNet ResNet and a 29 | [test image](https://github.com/BVLC/caffe/blob/master/examples/images/cat.jpg) 30 | 2. Converts the ResNet to an NNVM graph + library 31 | 3. Links the graph JSON definition, params, and runtime library into into an SGX 32 | enclave along with some code that performs inference. 33 | 4. Compiles and runs an executable that loads the enclave and requests that it perform 34 | inference on the image. 35 | which invokes the TVM module. 36 | 37 | For more information on building, please refer to the `Makefile`. 38 | For more information on the TVM module, please refer to `../howto_deploy`. 39 | For more in formation on SGX enclaves, please refer to the [SGX Enclave Demo](https://github.com/intel/linux-sgx/tree/master/SampleCode/SampleEnclave/) 40 | -------------------------------------------------------------------------------- /examples/sgx/build_model.py: -------------------------------------------------------------------------------- 1 | """Creates a neural network graph module, the system library, and params. 2 | Heavily inspired by tutorials/from_mxnet.py 3 | """ 4 | from __future__ import print_function 5 | import ast 6 | import os 7 | from os import path as osp 8 | import tempfile 9 | 10 | import mxnet as mx 11 | from mxnet.gluon.model_zoo.vision import get_model 12 | from mxnet.gluon.utils import download 13 | import nnvm 14 | import nnvm.compiler 15 | import numpy as np 16 | from PIL import Image 17 | import tvm 18 | 19 | 20 | EXAMPLE_ROOT = osp.abspath(osp.join(osp.dirname(__file__))) 21 | BIN_DIR = osp.join(EXAMPLE_ROOT, 'bin') 22 | LIB_DIR = osp.join(EXAMPLE_ROOT, 'lib') 23 | 24 | TVM_TARGET = 'llvm --system-lib' 25 | 26 | 27 | def _download_model_and_image(out_dir): 28 | mx_model = get_model('resnet18_v1', pretrained=True) 29 | 30 | img_path = osp.join(out_dir, 'cat.png') 31 | bin_img_path = osp.join(out_dir, 'cat.bin') 32 | download( 33 | 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true', 34 | img_path) 35 | img = Image.open(img_path).resize((224, 224)) 36 | img = _transform_image(img) 37 | img.astype('float32').tofile(bin_img_path) 38 | shape_dict = {'data': img.shape} 39 | 40 | return mx_model, shape_dict 41 | 42 | 43 | def _transform_image(image): 44 | image = np.array(image) - np.array([123., 117., 104.]) 45 | image /= np.array([58.395, 57.12, 57.375]) 46 | image = image.transpose((2, 0, 1)) 47 | image = image[np.newaxis, :] 48 | return image 49 | 50 | 51 | def main(): 52 | # load the model, input image, and imagenet classes 53 | mx_model, shape_dict = _download_model_and_image(BIN_DIR) 54 | 55 | # convert the model, add a softmax 56 | sym, params = nnvm.frontend.from_mxnet(mx_model) 57 | sym = nnvm.sym.softmax(sym) 58 | 59 | # build the graph 60 | graph, lib, params = nnvm.compiler.build( 61 | sym, TVM_TARGET, shape_dict, params=params) 62 | 63 | # save the built graph 64 | if not osp.isdir(LIB_DIR): 65 | os.mkdir(LIB_DIR) 66 | lib.save(osp.join(LIB_DIR, 'deploy_lib.o')) 67 | with open(osp.join(LIB_DIR, 'deploy_graph.json'), 'w') as f_graph_json: 68 | f_graph_json.write(graph.json()) 69 | with open(osp.join(LIB_DIR, 'deploy_params.bin'), 'wb') as f_params: 70 | f_params.write(nnvm.compiler.save_param_dict(params)) 71 | 72 | 73 | if __name__ == '__main__': 74 | main() 75 | -------------------------------------------------------------------------------- /examples/sgx/enclave.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "tvm/src/runtime/graph/graph_runtime.cc" 8 | #ifndef _LIBCPP_SGX_CONFIG 9 | #include 10 | #include 11 | #include 12 | #include "tvm/src/runtime/file_util.cc" 13 | #endif 14 | 15 | // the statically linked graph json and params 16 | extern char _binary_lib_deploy_params_bin_start[]; 17 | extern char _binary_lib_deploy_params_bin_end[]; 18 | extern char _binary_lib_deploy_graph_json_start[]; 19 | extern char _binary_lib_deploy_graph_json_end[]; 20 | 21 | int RunInference(const char* img) { 22 | tvm::runtime::Module graph_lib = 23 | (*tvm::runtime::Registry::Get("module._GetSystemLib"))(); 24 | 25 | size_t graph_json_size = ((size_t)_binary_lib_deploy_graph_json_end - 26 | (size_t)_binary_lib_deploy_graph_json_start); 27 | size_t graph_params_size = ((size_t)_binary_lib_deploy_params_bin_end - 28 | (size_t)_binary_lib_deploy_params_bin_start); 29 | std::string graph_json(_binary_lib_deploy_graph_json_start, graph_json_size); 30 | std::string graph_params(_binary_lib_deploy_params_bin_start, graph_params_size); 31 | 32 | int device_type = kDLCPU; 33 | int device_id = 0; 34 | 35 | TVMContext ctx; 36 | ctx.device_type = static_cast(device_type); 37 | ctx.device_id = device_id; 38 | std::shared_ptr graph_rt = 39 | std::make_shared(); 40 | 41 | graph_rt->Init(graph_json, graph_lib, ctx); 42 | graph_rt->LoadParams(graph_params); 43 | 44 | DLTensor* input; 45 | DLTensor* output; 46 | int ndim = 2; 47 | int dtype_code = kDLFloat; 48 | int dtype_bits = 32; 49 | int dtype_lanes = 1; 50 | 51 | int batch_size = 1; 52 | int64_t input_shape[4] = {batch_size, 3, 224, 224}; 53 | int64_t output_shape[1] = {1000 /* num_classes */}; 54 | TVMArrayAlloc(input_shape, 4 /* ndim */, dtype_code, dtype_bits, dtype_lanes, 55 | device_type, device_id, &input); 56 | TVMArrayAlloc(output_shape, 1, dtype_code, dtype_bits, dtype_lanes, 57 | device_type, device_id, &output); 58 | memcpy(input->data, img, sizeof(float)*batch_size*3*224*224); 59 | 60 | graph_rt->SetInput(graph_rt->GetInputIndex("data"), input); 61 | graph_rt->Run(); 62 | graph_rt->GetOutput(0, output); 63 | 64 | float max_prob = 0; 65 | unsigned max_class = -1; 66 | for (int i = 0; i < 1000; ++i) { 67 | float p = static_cast(output->data)[i]; 68 | if (p > max_prob) { 69 | max_prob = p; 70 | max_class = i; 71 | } 72 | } 73 | 74 | return max_class; 75 | } 76 | 77 | 78 | extern "C" { 79 | int ecall_infer(const char* img) { 80 | return RunInference(img); 81 | } 82 | } 83 | 84 | #ifndef _LIBCPP_SGX_CONFIG 85 | int main(void) { 86 | std::ifstream f_img("bin/cat.bin", std::ios::binary); 87 | std::string img(static_cast( 88 | std::stringstream() << f_img.rdbuf()).str()); 89 | unsigned predicted_class = RunInference(img.c_str()); 90 | if (predicted_class == 281) { 91 | std::cout << "It's a tabby!" << std::endl; 92 | return 0; 93 | } 94 | std::cerr << "Inference failed! Predicted class: " << 95 | predicted_class << std::endl; 96 | return -1; 97 | } 98 | #endif 99 | -------------------------------------------------------------------------------- /examples/sgx/enclave_config.xml: -------------------------------------------------------------------------------- 1 | 2 | 0 3 | 0 4 | 0x2000 5 | 0x6794000 6 | 1 7 | 1 8 | 0 9 | 0 10 | 0xFFFFFFFF 11 | 12 | -------------------------------------------------------------------------------- /examples/sgx/enclave_private.pem: -------------------------------------------------------------------------------- 1 | -----BEGIN RSA PRIVATE KEY----- 2 | MIIG4gIBAAKCAYEAroOogvsj/fZDZY8XFdkl6dJmky0lRvnWMmpeH41Bla6U1qLZ 3 | AmZuyIF+mQC/cgojIsrBMzBxb1kKqzATF4+XwPwgKz7fmiddmHyYz2WDJfAjIveJ 4 | ZjdMjM4+EytGlkkJ52T8V8ds0/L2qKexJ+NBLxkeQLfV8n1mIk7zX7jguwbCG1Pr 5 | nEMdJ3Sew20vnje+RsngAzdPChoJpVsWi/K7cettX/tbnre1DL02GXc5qJoQYk7b 6 | 3zkmhz31TgFrd9VVtmUGyFXAysuSAb3EN+5VnHGr0xKkeg8utErea2FNtNIgua8H 7 | ONfm9Eiyaav1SVKzPHlyqLtcdxH3I8Wg7yqMsaprZ1n5A1v/levxnL8+It02KseD 8 | 5HqV4rf/cImSlCt3lpRg8U5E1pyFQ2IVEC/XTDMiI3c+AR+w2jSRB3Bwn9zJtFlW 9 | KHG3m1xGI4ck+Lci1JvWWLXQagQSPtZTsubxTQNx1gsgZhgv1JHVZMdbVlAbbRMC 10 | 1nSuJNl7KPAS/VfzAgEDAoIBgHRXxaynbVP5gkO0ug6Qw/E27wzIw4SmjsxG6Wpe 11 | K7kfDeRskKxESdsA/xCrKkwGwhcx1iIgS5+Qscd1Yg+1D9X9asd/P7waPmWoZd+Z 12 | AhlKwhdPsO7PiF3e1AzHhGQwsUTt/Y/aSI1MpHBvy2/s1h9mFCslOUxTmWw0oj/Q 13 | ldIEgWeNR72CE2+jFIJIyml6ftnb6qzPiga8Bm48ubKh0kvySOqnkmnPzgh+JBD6 14 | JnBmtZbfPT97bwTT+N6rnPqOOApvfHPf15kWI8yDbprG1l4OCUaIUH1AszxLd826 15 | 5IPM+8gINLRDP1MA6azECPjTyHXhtnSIBZCyWSVkc05vYmNXYUNiXWMajcxW9M02 16 | wKzFELO8NCEAkaTPxwo4SCyIjUxiK1LbQ9h8PSy4c1+gGP4LAMR8xqP4QKg6zdu9 17 | osUGG/xRe/uufgTBFkcjqBHtK5L5VI0jeNIUAgW/6iNbYXjBMJ0GfauLs+g1VsOm 18 | WfdgXzsb9DYdMa0OXXHypmV4GwKBwQDUwQj8RKJ6c8cT4vcWCoJvJF00+RFL+P3i 19 | Gx2DLERxRrDa8AVGfqaCjsR+3vLgG8V/py+z+dxZYSqeB80Qeo6PDITcRKoeAYh9 20 | xlT3LJOS+k1cJcEmlbbO2IjLkTmzSwa80fWexKu8/Xv6vv15gpqYl1ngYoqJM3pd 21 | vzmTIOi7MKSZ0WmEQavrZj8zK4endE3v0eAEeQ55j1GImbypSf7Idh7wOXtjZ7WD 22 | Dg6yWDrri+AP/L3gClMj8wsAxMV4ZR8CgcEA0fzDHkFa6raVOxWnObmRoDhAtE0a 23 | cjUj976NM5yyfdf2MrKy4/RhdTiPZ6b08/lBC/+xRfV3xKVGzacm6QjqjZrUpgHC 24 | 0LKiZaMtccCJjLtPwQd0jGQEnKfMFaPsnhOc5y8qVkCzVOSthY5qhz0XNotHHFmJ 25 | gffVgB0iqrMTvSL7IA2yqqpOqNRlhaYhNl8TiFP3gIeMtVa9rZy31JPgT2uJ+kfo 26 | gV7sdTPEjPWZd7OshGxWpT6QfVDj/T9T7L6tAoHBAI3WBf2DFvxNL2KXT2QHAZ9t 27 | k3imC4f7U+wSE6zILaDZyzygA4RUbwG0gv8/TJVn2P/Eynf76DuWHGlaiLWnCbSz 28 | Az2DHBQBBaku409zDQym3j1ugMRjzzSQWzJg0SIyBH3hTmnYcn3+Uqcp/lEBvGW6 29 | O+rsXFt3pukqJmIV8HzLGGaLm62BHUeZf3dyWm+i3p/hQAL7Xvu04QW70xuGqdr5 30 | afV7p5eaeQIJXyGQJ0eylV/90+qxjMKiB1XYg6WYvwKBwQCL/ddpgOdHJGN8uRom 31 | e7Zq0Csi3hGheMKlKbN3vcxT5U7MdyHtTZZOJbTvxKNNUNYH/8uD+PqDGNneb29G 32 | BfGzvI3EASyLIcGZF3OhKwZd0jUrWk2y7Vhob91jwp2+t73vdMbkKyI4mHOuXvGv 33 | fg95si9oO7EBT+Oqvhccd2J+F1IVXncccYnF4u5ZGWt5lLewN/pVr7MjjykeaHqN 34 | t+rfnQam2psA6fL4zS2zTmZPzR2tnY8Y1GBTi0Ko1OKd1HMCgcAb5cB/7/AQlhP9 35 | yQa04PLH9ygQkKKptZp7dy5WcWRx0K/hAHRoi2aw1wZqfm7VBNu2SLcs90kCCCxp 36 | 6C5sfJi6b8NpNbIPC+sc9wsFr7pGo9SFzQ78UlcWYK2Gu2FxlMjonhka5hvo4zvg 37 | WxlpXKEkaFt3gLd92m/dMqBrHfafH7VwOJY2zT3WIpjwuk0ZzmRg5p0pG/svVQEH 38 | NZmwRwlopysbR69B/n1nefJ84UO50fLh5s5Zr3gBRwbWNZyzhXk= 39 | -----END RSA PRIVATE KEY----- 40 | -------------------------------------------------------------------------------- /examples/sgx/model.edl: -------------------------------------------------------------------------------- 1 | enclave { 2 | from "sgx_tstdc.edl" import sgx_thread_wait_untrusted_event_ocall, sgx_thread_set_untrusted_event_ocall, sgx_thread_setwait_untrusted_events_ocall, sgx_thread_set_multiple_untrusted_events_ocall; 3 | 4 | trusted { 5 | public unsigned ecall_infer([user_check] const char* img); 6 | }; 7 | }; 8 | -------------------------------------------------------------------------------- /examples/sgx/run_example.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | sgx_sdk=${SGX_SDK:=/opt/sgxsdk} 4 | 5 | make 6 | echo "=========================" 7 | LD_LIBRARY_PATH="$sgx_sdk/lib64":${LD_LIBRARY_PATH} bin/run_model 8 | -------------------------------------------------------------------------------- /include/nnvm/base.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file base.h 4 | * \brief Configuration of nnvm as well as basic data structure. 5 | */ 6 | #ifndef NNVM_BASE_H_ 7 | #define NNVM_BASE_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | namespace nnvm { 18 | 19 | /*! \brief any type */ 20 | using dmlc::any; 21 | 22 | /*! \brief array_veiw type */ 23 | using dmlc::array_view; 24 | 25 | /*!\brief getter function of any type */ 26 | using dmlc::get; 27 | 28 | } // namespace nnvm 29 | 30 | // describe op registration point 31 | #define NNVM_STRINGIZE_DETAIL(x) #x 32 | #define NNVM_STRINGIZE(x) NNVM_STRINGIZE_DETAIL(x) 33 | #define NNVM_DESCRIBE(...) describe(__VA_ARGS__ "\n\nFrom:" __FILE__ ":" NNVM_STRINGIZE(__LINE__)) 34 | #define NNVM_ADD_FILELINE "\n\nDefined in " __FILE__ ":L" NNVM_STRINGIZE(__LINE__) 35 | #endif // NNVM_BASE_H_ 36 | -------------------------------------------------------------------------------- /include/nnvm/compiler/op_attr_types.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file nnvm/compiler/op_attr_types.h 4 | * \brief The Expr and related elements in DataFlow construction. 5 | */ 6 | #ifndef NNVM_COMPILER_OP_ATTR_TYPES_H_ 7 | #define NNVM_COMPILER_OP_ATTR_TYPES_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include "packed_func_ext.h" 20 | 21 | namespace nnvm { 22 | namespace compiler { 23 | 24 | using ::tvm::Array; 25 | using ::tvm::Tensor; 26 | using ::tvm::Schedule; 27 | 28 | /*! \brief operator pattern used in graph fusion */ 29 | enum OpPatternKind { 30 | // Elementwise operation 31 | kElemWise = 0, 32 | // Broadcasting operator, can always map output axis to the input in order. 33 | // for example :code:`out[i, ax1, j, ax2] = input[i, j]`. 34 | // Note that the axis need to be in order so transpose is not a bcast operator. 35 | kBroadcast = 1, 36 | // Injective operator, can always injectively map output axis to a single input axis. 37 | // All injective operator can still be safely fused to injective and reduction. 38 | kInjective = 2, 39 | // Communicative reduction operator. 40 | kCommReduce = 3, 41 | // Complex operation, can still fuse elemwise operations into its output. 42 | // but cannot chain another complex op 43 | kOutEWiseFusable = 4, 44 | // Opaque operation, cannot fuse anything. 45 | kOpaque = 8 46 | }; 47 | 48 | /*! \brief the operator pattern */ 49 | using TOpPattern = int; 50 | 51 | /*! 52 | * \brief Computation description interface 53 | * \param attrs The attribute of the node. 54 | * \param inputs The input tensors(placeholders) 55 | * \param out_info Tensors holding shape/type information about output, 56 | & these are always placeholders. 57 | * \return The output description of the tensor. 58 | */ 59 | using FTVMCompute = std::function< 60 | Array(const NodeAttrs& attrs, 61 | const Array& inputs, 62 | const Array& out_info)>; 63 | 64 | /*! 65 | * \brief Build the computation schedule for 66 | * op whose root is at current op. 67 | * \param attrs The attribute of the node. 68 | * \param outs The output tensors. 69 | * \param target The build target. 70 | * \return schedule The computation schedule. 71 | */ 72 | using FTVMSchedule = std::function< 73 | Schedule(const NodeAttrs& attrs, 74 | const Array& outs, 75 | const std::string& target)>; 76 | 77 | /*! 78 | * \brief Modify the op node to alter its input layout. 79 | * it is invoked in AlterOpLayout pass. 80 | * \param attrs The attribute of the original node. 81 | * \param inputs The input symbols of the original node. 82 | * \param tinfos The inferred shape and dtype of the inputs. 83 | * \param ret The replaced operator. 84 | * \return Whether to replace current operator. 85 | */ 86 | using FTVMAlterOpLayout = std::function< 87 | bool(const NodeAttrs& attrs, 88 | const Symbol& inputs, 89 | const Array& tinfos, 90 | Symbol* ret)>; 91 | 92 | /*! 93 | * \brief Transform from normal operator to vectorized operator 94 | * \param node The source node. 95 | * \return Transformed vectorized op. 96 | */ 97 | using FTVMVectorizedOp = std::function; 98 | 99 | } // namespace compiler 100 | } // namespace nnvm 101 | #endif // NNVM_COMPILER_OP_ATTR_TYPES_H_ 102 | -------------------------------------------------------------------------------- /include/nnvm/compiler/packed_func_ext.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file packed_func_ext.h 4 | * \brief Extension to enable packed functionn for nnvm types 5 | */ 6 | #ifndef NNVM_COMPILER_PACKED_FUNC_EXT_H_ 7 | #define NNVM_COMPILER_PACKED_FUNC_EXT_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | namespace nnvm { 18 | namespace compiler { 19 | 20 | using tvm::runtime::PackedFunc; 21 | 22 | using AttrDict = std::unordered_map; 23 | 24 | /*! 25 | * \brief Get PackedFunction from global registry and 26 | * report error if it does not exist 27 | * \param name The name of the function. 28 | * \return The created PackedFunc. 29 | */ 30 | inline const PackedFunc& GetPackedFunc(const std::string& name) { 31 | const PackedFunc* pf = tvm::runtime::Registry::Get(name); 32 | CHECK(pf != nullptr) << "Cannot find function " << name << " in registry"; 33 | return *pf; 34 | } 35 | } // namespace compiler 36 | } // namespace nnvm 37 | 38 | // Enable the graph and symbol object exchange. 39 | namespace tvm { 40 | namespace runtime { 41 | 42 | template<> 43 | struct extension_class_info { 44 | static const int code = 16; 45 | }; 46 | 47 | template<> 48 | struct extension_class_info { 49 | static const int code = 17; 50 | }; 51 | 52 | template<> 53 | struct extension_class_info { 54 | static const int code = 18; 55 | }; 56 | 57 | } // namespace runtime 58 | } // namespace tvm 59 | #endif // NNVM_COMPILER_PACKED_FUNC_EXT_H_ 60 | -------------------------------------------------------------------------------- /include/nnvm/compiler/util.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file util.h 4 | * \brief Utility functions for nnvm compiler 5 | */ 6 | #ifndef NNVM_COMPILER_UTIL_H_ 7 | #define NNVM_COMPILER_UTIL_H_ 8 | 9 | #include 10 | #include 11 | 12 | namespace nnvm { 13 | namespace compiler { 14 | 15 | /* 16 | * \brief Helper function to convert TShape to TVM array. Useful for 17 | * passing data from NNVM param structures to TOPI ops. 18 | * 19 | * \param shape The shape to convert 20 | * 21 | * \return An Array of Expr, where each element is a constant int32 22 | */ 23 | inline tvm::Array ShapeToArray(TShape shape) { 24 | tvm::Array result; 25 | for (auto i : shape) { 26 | result.push_back(tvm::make_const(tvm::Int(32), i)); 27 | } 28 | return result; 29 | } 30 | 31 | } // namespace compiler 32 | } // namespace nnvm 33 | #endif // NNVM_COMPILER_UTIL_H_ 34 | -------------------------------------------------------------------------------- /include/nnvm/graph_attr_types.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file graph_attr_types.h 4 | * \brief Data structures that can appear in graph attributes. 5 | */ 6 | #ifndef NNVM_GRAPH_ATTR_TYPES_H_ 7 | #define NNVM_GRAPH_ATTR_TYPES_H_ 8 | 9 | #include 10 | #include 11 | #include "./tuple.h" 12 | #include "./layout.h" 13 | 14 | namespace nnvm { 15 | 16 | /*! 17 | * \brief The result holder of JSON serializer 18 | * 19 | * \note Stored under ret.attrs["json"], provided by Pass "SaveJSON" 20 | 21 | * \code 22 | * Graph ret = ApplyPass(src_graph, "SaveJSON"); 23 | * const JSONString& json = ret.GetAttr("shape"); 24 | * \endcode 25 | */ 26 | using JSONString = std::string; 27 | 28 | /*! 29 | * \brief The result holder of shape of each NodeEntry in the graph. 30 | * \note Stored under graph.attrs["shape"], provided by Pass "InferShape" 31 | * 32 | * \code 33 | * Graph g = ApplyPass(src_graph, "InferShape"); 34 | * const ShapeVector& shapes = g.GetAttr("shape"); 35 | * // get shape by entry id 36 | * TShape entry_shape = shapes[g.indexed_graph().entry_id(my_entry)]; 37 | * \endcode 38 | * 39 | * \sa FInferShape 40 | */ 41 | using ShapeVector = std::vector; 42 | 43 | /*! 44 | * \brief The result holder of type of each NodeEntry in the graph. 45 | * \note Stored under graph.attrs["dtype"], provided by Pass "InferType" 46 | * 47 | * \code 48 | * Graph g = ApplyPass(src_graph, "InferType"); 49 | * const DTypeVector& types = g.GetAttr("dtype"); 50 | * // get type by entry id 51 | * int entry_type = dtypes[g.indexed_graph().entry_id(my_entry)]; 52 | * \endcode 53 | * 54 | * \sa FInferType 55 | */ 56 | using DTypeVector = std::vector; 57 | 58 | /*! 59 | * \brief The result holder of layout of each NodeEntry in the graph. 60 | * \note Stored under graph.attrs["layout"], provided by Pass "InferType" 61 | * 62 | * \code 63 | * Graph g = ApplyPass(src_graph, "LayoutTransform"); 64 | * const LayoutVector& layouts = g.GetAttr("layout"); 65 | * // get layout by entry id 66 | * int entry_layout = layouts[g.indexed_graph().entry_id(my_entry)]; 67 | * \endcode 68 | * 69 | * \sa FCorrectLayout 70 | */ 71 | using LayoutVector = std::vector; 72 | 73 | /*! 74 | * \brief The result holder of device of each operator in the graph. 75 | * \note Stored under graph.attrs["device"], provided by Pass "PlaceDevice" 76 | * 77 | * \code 78 | * Graph g = ApplyPass(src_graph, "PlaceDevice"); 79 | * const &device = g.GetAttr("device"); 80 | * // get device by node_id 81 | * int device_type = device[g.indexed_graph().node_id(my_node)]; 82 | * \endcode 83 | */ 84 | using DeviceVector = std::vector; 85 | 86 | /*! 87 | * \brief The result holder of device of each operator in the graph. 88 | * 89 | * \note Stored under graph.attrs["device_assign_map"], needed by Pass "PlaceDevice" 90 | * -1 means unknown device 91 | */ 92 | using DeviceAssignMap = std::unordered_map; 93 | 94 | /*! 95 | * \brief The result holder of storage id of each NodeEntry in the graph. 96 | * 97 | * \note Stored under graph.attrs["storage"], provided by Pass "PlanMemory" 98 | * Storage id is a continuous integer. 99 | * If the storage id is -1 then the storage is not assigned. 100 | * 101 | * \code 102 | * Graph g = ApplyPass(src_graph, "PlanMemory"); 103 | * const &storage = g.GetAttr("storage"); 104 | * // get storage id by entry 105 | * int storage_id = storage[g.indexed_graph().entry_id(my_entry)]; 106 | * \endcode 107 | */ 108 | using StorageVector = std::vector; 109 | 110 | } // namespace nnvm 111 | 112 | #endif // NNVM_GRAPH_ATTR_TYPES_H_ 113 | -------------------------------------------------------------------------------- /include/nnvm/top/README: -------------------------------------------------------------------------------- 1 | NNVM Core Operator and Compiler 2 | -------------------------------------------------------------------------------- /make/config.mk: -------------------------------------------------------------------------------- 1 | #------------------------------------------------------------------------------- 2 | # Template configuration for compiling nnvm 3 | # 4 | # If you want to change the configuration, please use the following 5 | # steps. Assume you are on the root directory of nnvm. First copy the this 6 | # file so that any local changes will be ignored by git 7 | # 8 | # $ cp make/config.mk . 9 | # 10 | # Next modify the according entries, and then compile by 11 | # 12 | # $ make 13 | # 14 | # or build in parallel with 8 threads 15 | # 16 | # $ make -j8 17 | #------------------------------------------------------------------------------- 18 | 19 | #--------------------- 20 | # choice of compiler 21 | #-------------------- 22 | 23 | export NVCC = nvcc 24 | 25 | # choice of archiver 26 | export AR = ar 27 | 28 | # the additional link flags you want to add 29 | ADD_LDFLAGS= 30 | 31 | # the additional compile flags you want to add 32 | ADD_CFLAGS= 33 | 34 | # path to dmlc-core module 35 | #DMLC_CORE_PATH= 36 | 37 | #---------------------------- 38 | # plugins 39 | #---------------------------- 40 | 41 | # whether to use fusion integration. This requires installing cuda. 42 | # ifndef CUDA_PATH 43 | # CUDA_PATH = /usr/local/cuda 44 | # endif 45 | # NNVM_FUSION_PATH = plugin/nnvm-fusion 46 | # NNVM_PLUGINS += $(NNVM_FUSION_PATH)/nnvm-fusion.mk 47 | -------------------------------------------------------------------------------- /python/.gitignore: -------------------------------------------------------------------------------- 1 | *.c 2 | *.cpp 3 | -------------------------------------------------------------------------------- /python/conda/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | if [ -z "$PREFIX" ]; then 6 | PREFIX="$CONDA_PREFIX" 7 | fi 8 | 9 | if [ "$(uname)" = 'Darwin' ] 10 | then 11 | # Without this, Apple's default shipped clang will refuse to see any 12 | # headers like mutex. 13 | export MACOSX_DEPLOYMENT_TARGET=10.9 14 | fi 15 | 16 | rm -rf build || true 17 | mkdir -p build 18 | cd build 19 | # Enable static-libstdc++ to make it easier to link this library with 20 | # other C++ compilers 21 | CXXFLAGS=-static-libstdc++ cmake -DCMAKE_PREFIX_PATH=${PREFIX} -DCMAKE_INSTALL_PREFIX=${PREFIX} .. 22 | make -j4 VERBOSE=1 23 | make install/fast 24 | cd .. 25 | 26 | cd python 27 | $PYTHON setup.py install 28 | cd .. 29 | -------------------------------------------------------------------------------- /python/conda/meta.yaml: -------------------------------------------------------------------------------- 1 | {% set version = "0.1.dev" %} 2 | 3 | package: 4 | name: nnvm 5 | version: {{ version }} 6 | 7 | source: 8 | path: ../.. 9 | 10 | build: 11 | number: 1 12 | skip: True # [win] 13 | 14 | requirements: 15 | build: 16 | - cmake 17 | - python >=3 18 | - numpy 19 | - setuptools 20 | - nose 21 | - decorator 22 | run: 23 | - tvm 24 | - topi 25 | - python >=3 26 | - numpy 27 | - decorator 28 | 29 | about: 30 | home: https://github.com/dmlc/nnvm 31 | license: Apache2 32 | summary: Bring deep learning to bare metal 33 | -------------------------------------------------------------------------------- /python/nnvm/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # coding: utf-8 3 | """NNVM python API for ease of use and help new framework establish python API. """ 4 | from __future__ import absolute_import as _abs 5 | 6 | from . import _base 7 | from . import symbol as sym 8 | from . import symbol 9 | from ._base import NNVMError 10 | from . import frontend 11 | 12 | __version__ = _base.__version__ 13 | -------------------------------------------------------------------------------- /python/nnvm/_ctypes/README: -------------------------------------------------------------------------------- 1 | Ctypes specific implementation of certain modules -------------------------------------------------------------------------------- /python/nnvm/_ctypes/__init__.py: -------------------------------------------------------------------------------- 1 | """"ctypes implementation of the Symbol""" 2 | -------------------------------------------------------------------------------- /python/nnvm/_cy2/README: -------------------------------------------------------------------------------- 1 | This folder is by default empty and will hold DLLs generated by cython. 2 | -------------------------------------------------------------------------------- /python/nnvm/_cy2/__init__.py: -------------------------------------------------------------------------------- 1 | """Namespace for cython generated modules for python2""" 2 | -------------------------------------------------------------------------------- /python/nnvm/_cy3/README: -------------------------------------------------------------------------------- 1 | This folder is by default empty and will hold DLLs generated by cython. -------------------------------------------------------------------------------- /python/nnvm/_cy3/__init__.py: -------------------------------------------------------------------------------- 1 | """Cython generated modules""" 2 | -------------------------------------------------------------------------------- /python/nnvm/_symbol_internal.py: -------------------------------------------------------------------------------- 1 | """Module space to register internal functions. Leave empty""" 2 | -------------------------------------------------------------------------------- /python/nnvm/attribute.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | """Attribute scoping support for symbolic API.""" 3 | from __future__ import absolute_import 4 | 5 | from ._base import string_types 6 | 7 | class AttrScope(object): 8 | """Attribute manager for scoping. 9 | 10 | User can also inherit this object to change naming behavior. 11 | 12 | Parameters 13 | ---------- 14 | kwargs 15 | The attributes to set for all symbol creations in the scope. 16 | """ 17 | current = None 18 | 19 | def __init__(self, **kwargs): 20 | self._old_scope = None 21 | for value in kwargs.values(): 22 | if not isinstance(value, string_types): 23 | raise ValueError("Attributes need to be string") 24 | self._attr = kwargs 25 | 26 | def get(self, attr): 27 | """ 28 | Get the attribute dict given the attribute set by the symbol. 29 | 30 | Parameters 31 | ---------- 32 | attr : dict of string to string 33 | The attribute passed in by user during symbol creation. 34 | 35 | Returns 36 | ------- 37 | attr : dict of string to string 38 | Updated attributes to add other scope related attributes. 39 | """ 40 | if self._attr: 41 | ret = self._attr.copy() 42 | if attr: 43 | ret.update(attr) 44 | return ret 45 | else: 46 | return attr 47 | 48 | def __enter__(self): 49 | # pylint: disable=protected-access 50 | self._old_scope = AttrScope.current 51 | attr = AttrScope.current._attr.copy() 52 | attr.update(self._attr) 53 | self._attr = attr 54 | AttrScope.current = self 55 | return self 56 | 57 | def __exit__(self, ptype, value, trace): 58 | assert self._old_scope 59 | AttrScope.current = self._old_scope 60 | 61 | AttrScope.current = AttrScope() 62 | -------------------------------------------------------------------------------- /python/nnvm/compiler/__init__.py: -------------------------------------------------------------------------------- 1 | """NNVM compiler toolchain. 2 | 3 | User only need to use :any:`build` and :any:`build_config` to do the compilation, 4 | and :any:`save_param_dict` to save the parameters into bytes. 5 | The other APIs are for more advanced interaction with the compiler toolchain. 6 | """ 7 | from __future__ import absolute_import 8 | 9 | import tvm 10 | 11 | from . import build_module 12 | from . build_module import build, optimize, build_config 13 | from . compile_engine import engine, graph_key 14 | from . param_dict import save_param_dict, load_param_dict 15 | 16 | from .. import symbol as _symbol 17 | from .. import graph as _graph 18 | 19 | from .. import top as _top 20 | 21 | 22 | tvm.register_extension(_symbol.Symbol, _symbol.Symbol) 23 | tvm.register_extension(_graph.Graph, _graph.Graph) 24 | -------------------------------------------------------------------------------- /python/nnvm/compiler/compile_engine.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name 2 | """Compiler engine interface to internal engine 3 | 4 | You can get the engine singleton at ``nnvm.compiler.engine`` 5 | """ 6 | import tvm 7 | 8 | _list_cache_items = tvm.get_global_func("nnvm.compiler.ListCacheItems") 9 | _clear_cache = tvm.get_global_func("nnvm.compiler.ClearCache") 10 | _get_cache_item = tvm.get_global_func("nnvm.compiler.GetCacheItem") 11 | _set_cache_item = tvm.get_global_func("nnvm.compiler.SetCacheItem") 12 | _graph_key_get_graph = tvm.get_global_func("nnvm.compiler.GraphKeyGetGraph") 13 | _make_graph_key = tvm.get_global_func("nnvm.compiler.MakeGraphKey") 14 | 15 | @tvm.register_node 16 | class GraphKey(tvm.node.NodeBase): 17 | """Key of a graph compilation context""" 18 | @property 19 | def graph(self): 20 | return _graph_key_get_graph(self) 21 | 22 | 23 | @tvm.register_node 24 | class GraphCacheEntry(tvm.node.NodeBase): 25 | """CacheEntry of compilation into a TVM Function""" 26 | pass 27 | 28 | 29 | @tvm.register_node 30 | class GraphFunc(tvm.node.NodeBase): 31 | """Compiled result of a graph into a TVM Function""" 32 | pass 33 | 34 | 35 | class Engine(object): 36 | """Global singleton compilation engine. 37 | 38 | You can get the singleton at ``nnvm.compiler.engine`` 39 | """ 40 | def items(self): 41 | """List the available cache key value pairs. 42 | 43 | Returns 44 | ------- 45 | item_list : list of (GraphKey, GraphCacheEntry) 46 | The existing cache items 47 | """ 48 | res = _list_cache_items() 49 | assert len(res) % 2 == 0 50 | return [(res[2*i], res[2*i+1]) for i in range(len(res) // 2)] 51 | 52 | def clear_cache(self): 53 | """Clear the existing cached functions.""" 54 | _clear_cache() 55 | 56 | def __setitem__(self, key, value): 57 | """Clear the existing cached functions.""" 58 | if isinstance(value, GraphCacheEntry): 59 | _set_cache_item(key, value.graph_func) 60 | else: 61 | _set_cache_item(key, value) 62 | 63 | def __getitem__(self, key): 64 | """Clear the existing cached functions.""" 65 | return _get_cache_item(key) 66 | 67 | def dump(self): 68 | """Return a string representation of engine dump 69 | 70 | Returns 71 | ------- 72 | dump : str 73 | The dumped string representation 74 | """ 75 | items = self.items() 76 | res = "====================================\n" 77 | res += "CompilerEngine dump, %d items cached\n" % len(items) 78 | for key, value in items: 79 | res += "------------------------------------\n" 80 | res += "target={}\n".format(key.target) 81 | res += "inputs={}\n".format(key.inputs) 82 | res += "use_count={}\n".format(value.use_count) 83 | res += "func_name={}\n".format(value.graph_func.func_name) 84 | res += key.graph.ir() + "\n" 85 | res += "===================================\n" 86 | return res 87 | 88 | engine = Engine() 89 | 90 | 91 | def graph_key(graph, inputs, target): 92 | """Construct a new graph key. 93 | 94 | Parameters 95 | ---------- 96 | graph : Graph 97 | The computation graph structure 98 | 99 | inputs : list of Tensor(placeholder) 100 | The input requirement to the graph. 101 | 102 | target : str 103 | The target of compilation. 104 | """ 105 | return _make_graph_key(graph, inputs, target) 106 | -------------------------------------------------------------------------------- /python/nnvm/compiler/graph_attr.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name 2 | """Utilities to access graph attributes""" 3 | from __future__ import absolute_import as _abs 4 | 5 | import tvm 6 | 7 | def set_shape_inputs(g, shape): 8 | """Set the shape of input graph nodes in the graph attribute. 9 | 10 | Parameters 11 | ---------- 12 | g : Graph 13 | The input graph 14 | 15 | shape : dict of str to tuple 16 | The input shape 17 | 18 | Returns 19 | ------- 20 | g : Graph 21 | The updated graph with updated shape. 22 | """ 23 | list_shape = [ 24 | shape.get(name, ()) for name in g.index.input_names] 25 | g._set_json_attr("shape_inputs", list_shape, 'list_shape') 26 | return g 27 | 28 | 29 | DTYPE_TO_TCODE = { 30 | "default": -1, 31 | "float32": 0, 32 | "float64": 1, 33 | "float16": 2, 34 | "uint8": 3, 35 | "int32": 4, 36 | "int8": 5, 37 | "int64": 6, 38 | "int16": 7, 39 | "uint16": 8, 40 | "uint32": 9, 41 | "uint64": 10, 42 | } 43 | 44 | TCODE_TO_DTYPE = { 45 | -1: None, 46 | 0: "float32", 47 | 1: "float64", 48 | 2: "float16", 49 | 3: "uint8", 50 | 4: "int32", 51 | 5: "int8", 52 | 6: "int64", 53 | 7: "int16", 54 | 8: "uint16", 55 | 9: "uint32", 56 | 10: "uint64", 57 | } 58 | 59 | def set_dtype_inputs(g, dtype): 60 | """Set the dtype inputs of graph nodes 61 | 62 | Parameters 63 | ---------- 64 | g : Graph 65 | The input graph 66 | 67 | dtype : dict of str to str or str 68 | The input dtype 69 | 70 | Returns 71 | ------- 72 | g : Graph 73 | The updated graph with updated dtype. 74 | """ 75 | if isinstance(dtype, dict): 76 | list_dtype = [ 77 | DTYPE_TO_TCODE[str(dtype.get(name, "default"))] 78 | for name in g.index.input_names] 79 | else: 80 | list_dtype = [DTYPE_TO_TCODE[dtype]] * len(g.index.input_names) 81 | g._set_json_attr("dtype_inputs", list_dtype, "list_int") 82 | return g 83 | 84 | 85 | def set_layout_inputs(g, layout): 86 | """Set the layout inputs of graph nodes 87 | 88 | Parameters 89 | ---------- 90 | g : Graph 91 | The input graph 92 | 93 | layout : dict of str to str or str 94 | The input layout 95 | 96 | Returns 97 | ------- 98 | g : Graph 99 | The updated graph with updated layout. 100 | """ 101 | if isinstance(layout, dict): 102 | list_layout = [ 103 | layout.get(name, "__undef__") for name in g.index.input_names] 104 | elif isinstance(layout, str): 105 | list_layout = ["__undef__"] * len(g.index.input_names) 106 | list_layout[0] = layout 107 | else: 108 | raise ValueError("Input layout must be str or dict") 109 | last_inferred_layouts = g.json_attr("layout") 110 | if last_inferred_layouts: 111 | input_layout = [last_inferred_layouts[g.index.entry_id(x)] for x in g.index.input_names] 112 | for i, layout_stored in enumerate(input_layout): 113 | list_layout[i] = list_layout[i] if list_layout[i] != '__undef__' else layout_stored 114 | g._set_json_attr("layout_inputs", list_layout, 'list_layout') 115 | return g 116 | 117 | _move_out_module = tvm.get_global_func("nnvm.graph._move_module") 118 | _move_out_graph = tvm.get_global_func("nnvm.graph._move_graph") 119 | -------------------------------------------------------------------------------- /python/nnvm/compiler/graph_pass.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name 2 | """Namespace of graph pass. 3 | 4 | Principle: 5 | - Graph in, graph out: always takes in graph as first argument and returns a graph 6 | - Composable API: break graph transformation pass as segments of small transformations. 7 | """ 8 | from __future__ import absolute_import as _abs 9 | -------------------------------------------------------------------------------- /python/nnvm/compiler/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=too-few-public-methods, no-member 2 | """API for scheduling learning rate.""" 3 | from .. import symbol as sym 4 | 5 | class LRScheduler(object): 6 | """Base class of a learning rate scheduler. 7 | 8 | A scheduler returns a new learning rate based on the number of updates that have 9 | been performed. 10 | 11 | Parameters 12 | ---------- 13 | base_lr : float, optional 14 | The initial learning rate. 15 | """ 16 | def __init__(self, base_lr=0.01, name='LRScheduler'): 17 | self.name = name 18 | self.base_lr = base_lr 19 | 20 | def __call__(self, num_update): 21 | """Return a new learning rate based on number of updates. 22 | 23 | Parameters 24 | ---------- 25 | num_update: nnvm Symbol 26 | the number of updates applied to weight. 27 | """ 28 | raise NotImplementedError("__call__ method must be overridden.") 29 | 30 | class FactorScheduler(LRScheduler): 31 | """Reduce the learning rate by a factor for every *n* steps. 32 | 33 | It returns a new learning rate by:: 34 | 35 | base_lr * pow(factor, num_update/step) 36 | 37 | Parameters 38 | ---------- 39 | step : int 40 | Changes the learning rate for every n updates. 41 | factor : float, optional 42 | The factor to change the learning rate. 43 | stop_factor_lr : float, optional 44 | Stop updating the learning rate if it is less than this value. 45 | """ 46 | def __init__(self, step, factor=1, stop_factor_lr=1e-8, name='FactorScheduler', **kwargs): 47 | super(FactorScheduler, self).__init__(name=name, **kwargs) 48 | if step < 1: 49 | raise ValueError("Schedule step must be greater or equal than 1 round") 50 | if factor > 1.0: 51 | raise ValueError("Factor must be no more than 1 to make lr reduce") 52 | self.step = step 53 | self.factor = factor 54 | self.stop_factor_lr = stop_factor_lr 55 | 56 | def __call__(self, num_update): 57 | updated_lr = self.base_lr * self.factor ** (num_update / self.step) 58 | return sym.clip(updated_lr, a_min=self.stop_factor_lr, a_max=self.base_lr) 59 | -------------------------------------------------------------------------------- /python/nnvm/compiler/param_dict.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name 2 | """Helper utility to save parameter dict""" 3 | import ctypes 4 | import tvm 5 | from tvm._ffi.runtime_ctypes import TVMArrayHandle 6 | 7 | _save_param_dict = tvm.get_global_func("nnvm.compiler._save_param_dict") 8 | _load_param_dict = tvm.get_global_func("nnvm.compiler._load_param_dict") 9 | 10 | def save_param_dict(params): 11 | """Save parameter dictionary to binary bytes. 12 | 13 | The result binary bytes can be loaded by the 14 | GraphModule with API "load_params". 15 | 16 | Parameters 17 | ---------- 18 | params : dict of str to NDArray 19 | The parameter dictionary. 20 | 21 | Returns 22 | ------- 23 | param_bytes: bytearray 24 | Serialized parameters. 25 | 26 | Examples 27 | -------- 28 | .. code-block:: python 29 | 30 | # compile and save the modules to file. 31 | graph, lib, params = nnvm.compiler.build( 32 | graph, target, shape={"data", data_shape}, params=params) 33 | module = graph_runtime.create(graph, lib, tvm.gpu(0)) 34 | # save the parameters as byte array 35 | param_bytes = nnvm.compiler.save_param_dict(params) 36 | # We can serialize the param_bytes and load it back later. 37 | # Pass in byte array to module to directly set parameters 38 | module["load_params"](param_bytes) 39 | """ 40 | args = [] 41 | for k, v in params.items(): 42 | args.append(k) 43 | args.append(tvm.nd.array(v)) 44 | return _save_param_dict(*args) 45 | 46 | 47 | def load_param_dict(param_bytes): 48 | """Load parameter dictionary to binary bytes. 49 | 50 | Parameters 51 | ---------- 52 | param_bytes: bytearray 53 | Serialized parameters. 54 | 55 | Returns 56 | ------- 57 | params : dict of str to NDArray 58 | The parameter dictionary. 59 | """ 60 | if isinstance(param_bytes, (bytes, str)): 61 | param_bytes = bytearray(param_bytes) 62 | load_mod = _load_param_dict(param_bytes) 63 | size = load_mod(0) 64 | param_dict = {} 65 | for i in range(size): 66 | key = load_mod(1, i) 67 | dltensor_handle = ctypes.cast(load_mod(2, i), TVMArrayHandle) 68 | param_dict[key] = tvm.nd.NDArray(dltensor_handle, False) 69 | return param_dict 70 | -------------------------------------------------------------------------------- /python/nnvm/contrib.py: -------------------------------------------------------------------------------- 1 | """Module space to register contrib functions. Leave empty""" 2 | -------------------------------------------------------------------------------- /python/nnvm/cython/README: -------------------------------------------------------------------------------- 1 | Cython specific implementation of certain modules -------------------------------------------------------------------------------- /python/nnvm/cython/base.pyi: -------------------------------------------------------------------------------- 1 | ctypedef void* SymbolHandle 2 | ctypedef void* OpHandle 3 | ctypedef unsigned nn_uint 4 | 5 | cdef py_str(const char* x): 6 | if PY_MAJOR_VERSION < 3: 7 | return x 8 | else: 9 | return x.decode("utf-8") 10 | 11 | 12 | cdef c_str(pystr): 13 | """Create ctypes char * from a python string 14 | Parameters 15 | ---------- 16 | string : string type 17 | python string 18 | 19 | Returns 20 | ------- 21 | str : c_char_p 22 | A char pointer that can be passed to C API 23 | """ 24 | return pystr.encode("utf-8") 25 | 26 | 27 | cdef CALL(int ret): 28 | if ret != 0: 29 | raise NNVMError(NNGetLastError()) 30 | 31 | 32 | cdef const char** CBeginPtr(vector[const char*]& vec): 33 | if (vec.size() != 0): 34 | return &vec[0] 35 | else: 36 | return NULL 37 | 38 | cdef vector[const char*] SVec2Ptr(vector[string]& vec): 39 | cdef vector[const char*] svec 40 | svec.resize(vec.size()) 41 | for i in range(vec.size()): 42 | svec[i] = vec[i].c_str() 43 | return svec 44 | 45 | 46 | cdef BuildDoc(nn_uint num_args, 47 | const char** arg_names, 48 | const char** arg_types, 49 | const char** arg_descs, 50 | remove_dup=True): 51 | """Convert ctypes returned doc string information into parameters docstring. 52 | 53 | num_args : nn_uint 54 | Number of arguments. 55 | 56 | arg_names : ctypes.POINTER(ctypes.c_char_p) 57 | Argument names. 58 | 59 | arg_types : ctypes.POINTER(ctypes.c_char_p) 60 | Argument type information. 61 | 62 | arg_descs : ctypes.POINTER(ctypes.c_char_p) 63 | Argument description information. 64 | 65 | remove_dup : boolean, optional 66 | Whether remove duplication or not. 67 | 68 | Returns 69 | ------- 70 | docstr : str 71 | Python docstring of parameter sections. 72 | """ 73 | param_keys = set() 74 | param_str = [] 75 | for i in range(num_args): 76 | key = arg_names[i] 77 | if key in param_keys and remove_dup: 78 | continue 79 | param_keys.add(key) 80 | type_info = arg_types[i] 81 | ret = '%s : %s' % (key, type_info) 82 | if len(arg_descs[i]) != 0: 83 | ret += '\n ' + py_str(arg_descs[i]) 84 | param_str.append(ret) 85 | doc_str = ('Parameters\n' + 86 | '----------\n' + 87 | '%s\n') 88 | doc_str = doc_str % ('\n'.join(param_str)) 89 | return doc_str 90 | -------------------------------------------------------------------------------- /python/nnvm/frontend/__init__.py: -------------------------------------------------------------------------------- 1 | """NNVM frontends.""" 2 | from __future__ import absolute_import 3 | from .mxnet import from_mxnet 4 | from .onnx import from_onnx 5 | from .coreml import from_coreml 6 | from .keras import from_keras 7 | from .darknet import from_darknet 8 | -------------------------------------------------------------------------------- /python/nnvm/libinfo.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | """Information about nnvm.""" 3 | from __future__ import absolute_import 4 | import sys 5 | import os 6 | import platform 7 | 8 | if sys.version_info[0] == 3: 9 | import builtins as __builtin__ 10 | else: 11 | import __builtin__ 12 | 13 | def find_lib_path(): 14 | """Find NNNet dynamic library files. 15 | 16 | Returns 17 | ------- 18 | lib_path : list(string) 19 | List of all found path to the libraries 20 | """ 21 | if hasattr(__builtin__, "NNVM_BASE_PATH"): 22 | base_path = __builtin__.NNVM_BASE_PATH 23 | else: 24 | base_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) 25 | 26 | if hasattr(__builtin__, "NNVM_LIBRARY_NAME"): 27 | lib_name = __builtin__.NNVM_LIBRARY_NAME 28 | else: 29 | lib_name = "nnvm_compiler" if sys.platform.startswith('win32') else "libnnvm_compiler" 30 | 31 | api_path = os.path.join(base_path, '../../lib/') 32 | cmake_build_path = os.path.join(base_path, '../../build/Release/') 33 | cmake_build_path = os.path.join(base_path, '../../build/') 34 | dll_path = [base_path, api_path, cmake_build_path] 35 | 36 | if sys.platform.startswith('linux') and os.environ.get('LD_LIBRARY_PATH', None): 37 | dll_path.extend([p.strip() for p in os.environ['LD_LIBRARY_PATH'].split(":")]) 38 | elif sys.platform.startswith('darwin') and os.environ.get('DYLD_LIBRARY_PATH', None): 39 | dll_path.extend([p.strip() for p in os.environ['DYLD_LIBRARY_PATH'].split(":")]) 40 | elif sys.platform.startswith('win32') and os.environ.get('PATH', None): 41 | dll_path.extend([p.strip() for p in os.environ['PATH'].split(";")]) 42 | 43 | if sys.platform.startswith('win32'): 44 | vs_configuration = 'Release' 45 | if platform.architecture()[0] == '64bit': 46 | dll_path.append(os.path.join(base_path, '../../build', vs_configuration)) 47 | dll_path.append(os.path.join(base_path, '../../windows/x64', vs_configuration)) 48 | else: 49 | dll_path.append(os.path.join(base_path, '../../build', vs_configuration)) 50 | dll_path.append(os.path.join(base_path, '../../windows', vs_configuration)) 51 | dll_path = [os.path.join(p, '%s.dll' % lib_name) for p in dll_path] 52 | elif sys.platform.startswith('darwin'): 53 | dll_path = [os.path.join(p, '%s.dylib' % lib_name) for p in dll_path] 54 | else: 55 | dll_path = [os.path.join(p, '%s.so' % lib_name) for p in dll_path] 56 | 57 | lib_path = [p for p in dll_path if os.path.exists(p) and os.path.isfile(p)] 58 | if not lib_path: 59 | raise RuntimeError('Cannot find the files.\n' + 60 | 'List of candidates:\n' + str('\n'.join(dll_path))) 61 | return lib_path 62 | 63 | 64 | # current version 65 | __version__ = "0.8.0" 66 | -------------------------------------------------------------------------------- /python/nnvm/name.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | """Automatic naming support for symbolic API.""" 3 | from __future__ import absolute_import as _abs 4 | 5 | class NameManager(object): 6 | """NameManager to do automatic naming. 7 | 8 | User can also inherit this object to change naming behavior. 9 | """ 10 | current = None 11 | 12 | def __init__(self): 13 | self._counter = {} 14 | self._old_manager = None 15 | 16 | def get(self, name, hint): 17 | """Get the canonical name for a symbol. 18 | 19 | This is default implementation. 20 | When user specified a name, 21 | the user specified name will be used. 22 | 23 | When user did not, we will automatically generate a 24 | name based on hint string. 25 | 26 | Parameters 27 | ---------- 28 | name : str or None 29 | The name user specified. 30 | 31 | hint : str 32 | A hint string, which can be used to generate name. 33 | 34 | Returns 35 | ------- 36 | full_name : str 37 | A canonical name for the user. 38 | """ 39 | if name: 40 | return name 41 | if hint not in self._counter: 42 | self._counter[hint] = 0 43 | name = '%s%d' % (hint, self._counter[hint]) 44 | self._counter[hint] += 1 45 | return name 46 | 47 | def __enter__(self): 48 | self._old_manager = NameManager.current 49 | NameManager.current = self 50 | return self 51 | 52 | def __exit__(self, ptype, value, trace): 53 | assert self._old_manager 54 | NameManager.current = self._old_manager 55 | 56 | 57 | class Prefix(NameManager): 58 | """A name manager that always attach a prefix to all names. 59 | 60 | Examples 61 | -------- 62 | >>> import nnvm as nn 63 | >>> data = nn.symbol.Variable('data') 64 | >>> with nn.name.Prefix('mynet_'): 65 | net = nn.symbol.FullyConnected(data, num_hidden=10, name='fc1') 66 | >>> net.list_arguments() 67 | ['data', 'mynet_fc1_weight', 'mynet_fc1_bias'] 68 | """ 69 | def __init__(self, prefix): 70 | super(Prefix, self).__init__() 71 | self._prefix = prefix 72 | 73 | def get(self, name, hint): 74 | name = super(Prefix, self).get(name, hint) 75 | return self._prefix + name 76 | 77 | # initialize the default name manager 78 | NameManager.current = NameManager() 79 | -------------------------------------------------------------------------------- /python/nnvm/testing/__init__.py: -------------------------------------------------------------------------------- 1 | """Utilities for testing and benchmarks""" 2 | from __future__ import absolute_import as _abs 3 | 4 | from .config import ctx_list 5 | from .utils import create_workload 6 | from . import mobilenet 7 | from . import mlp 8 | from . import resnet 9 | from . import vgg 10 | from . import yolo2_detection 11 | -------------------------------------------------------------------------------- /python/nnvm/testing/config.py: -------------------------------------------------------------------------------- 1 | """Configuration about tests""" 2 | from __future__ import absolute_import as _abs 3 | 4 | import os 5 | import tvm 6 | 7 | def ctx_list(): 8 | """Get context list for testcases""" 9 | device_list = os.environ.get("NNVM_TEST_TARGETS", "") 10 | device_list = (device_list.split(",") if device_list 11 | else ["llvm", "cuda"]) 12 | device_list = set(device_list) 13 | res = [("llvm", tvm.cpu(0)), ("cuda", tvm.gpu(0))] 14 | return [x for x in res if x[1].exist and x[0] in device_list] 15 | -------------------------------------------------------------------------------- /python/nnvm/testing/init.py: -------------------------------------------------------------------------------- 1 | """Initializer of parameters.""" 2 | import numpy as np 3 | 4 | class Initializer(object): 5 | """The base class of an initializer.""" 6 | def __init__(self, **kwargs): 7 | self._kwargs = kwargs 8 | 9 | def __call__(self, desc, arr): 10 | """Initialize an array 11 | 12 | Parameters 13 | ---------- 14 | desc : str 15 | Initialization pattern descriptor. 16 | 17 | arr : NDArray 18 | The array to be initialized. 19 | """ 20 | if desc.endswith('weight'): 21 | self._init_weight(desc, arr) 22 | elif desc.endswith('bias'): 23 | self._init_bias(desc, arr) 24 | elif desc.endswith('gamma'): 25 | self._init_gamma(desc, arr) 26 | elif desc.endswith('beta'): 27 | self._init_beta(desc, arr) 28 | elif desc.endswith('mean'): 29 | self._init_mean(desc, arr) 30 | elif desc.endswith('var'): 31 | self._init_var(desc, arr) 32 | else: 33 | self._init_default(desc, arr) 34 | 35 | def _init_bias(self, _, arr): 36 | arr[:] = 0.0 37 | 38 | def _init_gamma(self, _, arr): 39 | arr[:] = 1.0 40 | 41 | def _init_beta(self, _, arr): 42 | arr[:] = 0.0 43 | 44 | def _init_mean(self, _, arr): 45 | arr[:] = 0.0 46 | 47 | def _init_var(self, _, arr): 48 | arr[:] = 1.0 49 | 50 | def _init_weight(self, name, arr): 51 | """Abstract method to Initialize weight.""" 52 | raise NotImplementedError("Must override it") 53 | 54 | def _init_default(self, name, _): 55 | raise ValueError( 56 | 'Unknown initialization pattern for %s. ' \ 57 | 'Default initialization is now limited to '\ 58 | '"weight", "bias", "gamma" (1.0), and "beta" (0.0).' \ 59 | 'Please use mx.sym.Variable(init=mx.init.*) to set initialization pattern' % name) 60 | 61 | 62 | class Xavier(Initializer): 63 | """ "Xavier" initialization for weights 64 | 65 | Parameters 66 | ---------- 67 | rnd_type: str, optional 68 | Random generator type, can be ``'gaussian'`` or ``'uniform'``. 69 | 70 | factor_type: str, optional 71 | Can be ``'avg'``, ``'in'``, or ``'out'``. 72 | 73 | magnitude: float, optional 74 | Scale of random number. 75 | """ 76 | def __init__(self, rnd_type="uniform", factor_type="avg", magnitude=3): 77 | super(Xavier, self).__init__(rnd_type=rnd_type, 78 | factor_type=factor_type, 79 | magnitude=magnitude) 80 | self.rnd_type = rnd_type 81 | self.factor_type = factor_type 82 | self.magnitude = float(magnitude) 83 | 84 | def _init_weight(self, name, arr): 85 | shape = arr.shape 86 | hw_scale = 1. 87 | if len(shape) < 2: 88 | raise ValueError('Xavier initializer cannot be applied to vector {0}. It requires at' 89 | ' least 2D.'.format(name)) 90 | if len(shape) > 2: 91 | hw_scale = np.prod(shape[2:]) 92 | fan_in, fan_out = shape[1] * hw_scale, shape[0] * hw_scale 93 | factor = 1. 94 | if self.factor_type == "avg": 95 | factor = (fan_in + fan_out) / 2.0 96 | elif self.factor_type == "in": 97 | factor = fan_in 98 | elif self.factor_type == "out": 99 | factor = fan_out 100 | else: 101 | raise ValueError("Incorrect factor type") 102 | # Hack for mobilenet, because there is less connectivity 103 | if "depthwise" in name: 104 | factor = 3 * 3 105 | scale = np.sqrt(self.magnitude / factor) 106 | if self.rnd_type == "uniform": 107 | arr[:] = np.random.uniform(-scale, scale, size=arr.shape) 108 | else: 109 | raise ValueError("Unknown random type") 110 | -------------------------------------------------------------------------------- /python/nnvm/testing/mlp.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | """ 18 | a simple multilayer perceptron 19 | """ 20 | from .. import symbol as sym 21 | from . utils import create_workload 22 | 23 | def get_symbol(num_classes=1000): 24 | data = sym.Variable('data') 25 | data = sym.flatten(data=data) 26 | fc1 = sym.dense(data=data, name='fc1', units=128) 27 | act1 = sym.relu(data=fc1, name='relu1') 28 | fc2 = sym.dense(data=act1, name='fc2', units=64) 29 | act2 = sym.relu(data=fc2, name='relu2') 30 | fc3 = sym.dense(data=act2, name='fc3', units=num_classes) 31 | mlp = sym.softmax(data=fc3, name='softmax') 32 | return mlp 33 | 34 | def get_workload(batch_size, num_classes=1000, image_shape=(3, 224, 224), dtype="float32"): 35 | """Get benchmark workload for a simple multilayer perceptron 36 | 37 | Parameters 38 | ---------- 39 | batch_size : int 40 | The batch size used in the model 41 | 42 | num_classes : int, optional 43 | Number of claseses 44 | 45 | image_shape : tuple, optional 46 | The input image shape 47 | 48 | dtype : str, optional 49 | The data type 50 | 51 | Returns 52 | ------- 53 | net : nnvm.symbol 54 | The computational graph 55 | 56 | params : dict of str to NDArray 57 | The parameters. 58 | """ 59 | net = get_symbol(num_classes=num_classes) 60 | return create_workload(net, batch_size, image_shape, dtype) 61 | -------------------------------------------------------------------------------- /python/nnvm/testing/utils.py: -------------------------------------------------------------------------------- 1 | """Helper utility to create common workload for testing.""" 2 | from __future__ import absolute_import as _abs 3 | 4 | import numpy as np 5 | import tvm 6 | from ..compiler import graph_util 7 | from ..import graph 8 | from . init import Xavier 9 | 10 | def create_workload(net, batch_size, image_shape=(3, 224, 224), 11 | dtype="float32", initializer=None, seed=0): 12 | """Helper function to create benchmark workload for input network 13 | 14 | Parameters 15 | ---------- 16 | net : nnvm.Symbol 17 | The selected network symbol to use 18 | 19 | batch_size : int 20 | The batch size used in the model 21 | 22 | image_shape : tuple, optional 23 | The input image shape 24 | 25 | dtype : str, optional 26 | The data type 27 | 28 | initializer : Initializer 29 | The initializer used 30 | 31 | seed : int 32 | The seed used in initialization. 33 | 34 | Returns 35 | ------- 36 | net : nnvm.Symbol 37 | The computational graph 38 | 39 | params : dict of str to NDArray 40 | The parameters. 41 | """ 42 | if image_shape is None: 43 | image_shape = (3, 224, 224) 44 | data_shape = (batch_size,) + image_shape 45 | params = {} 46 | g = graph.create(net) 47 | input_shapes, _ = graph_util.infer_shape(g, data=data_shape) 48 | shape_dict = dict(zip(g.index.input_names, input_shapes)) 49 | np.random.seed(seed) 50 | initializer = initializer if initializer else Xavier() 51 | for k, v in shape_dict.items(): 52 | if k == "data": 53 | continue 54 | init_value = np.zeros(v).astype(dtype) 55 | initializer(k, init_value) 56 | params[k] = tvm.nd.array(init_value, ctx=tvm.cpu(0)) 57 | return net, params 58 | -------------------------------------------------------------------------------- /python/nnvm/top/__init__.py: -------------------------------------------------------------------------------- 1 | """Tensor operator property registry 2 | 3 | Provide information to lower and schedule tensor operators. 4 | """ 5 | from .attr_dict import AttrDict 6 | from . import tensor 7 | from . import nn 8 | from . import transform 9 | from . import reduction 10 | from . import vision 11 | 12 | from .registry import OpPattern 13 | from .registry import register_compute, register_schedule, register_pattern 14 | -------------------------------------------------------------------------------- /python/nnvm/top/reduction.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name, unused-argument 2 | """Reduction ops""" 3 | from __future__ import absolute_import 4 | 5 | import tvm 6 | import topi 7 | import topi.cuda 8 | from . import registry as reg 9 | from .registry import OpPattern 10 | 11 | def _schedule_reduce(_, outs, target): 12 | """Generic schedule for reduce""" 13 | with tvm.target.create(target): 14 | return topi.generic.schedule_reduce(outs) 15 | 16 | 17 | _fschedule_reduce = tvm.convert(_schedule_reduce) 18 | 19 | def _compute_reduce(f): 20 | """auxiliary function""" 21 | def _compute(attrs, inputs, out_info): 22 | axis = attrs.get_int_tuple("axis") 23 | keepdims = attrs.get_bool("keepdims") 24 | if axis: 25 | return f(inputs[0], axis=axis, keepdims=keepdims) 26 | return f(inputs[0], keepdims=keepdims) 27 | return _compute 28 | 29 | # sum 30 | reg.register_pattern("sum", OpPattern.COMM_REDUCE) 31 | reg.register_schedule("sum", _fschedule_reduce) 32 | 33 | # max 34 | reg.register_pattern("max", OpPattern.COMM_REDUCE) 35 | reg.register_schedule("max", _fschedule_reduce) 36 | 37 | # min 38 | reg.register_pattern("min", OpPattern.COMM_REDUCE) 39 | reg.register_schedule("min", _fschedule_reduce) 40 | -------------------------------------------------------------------------------- /python/nnvm/top/registry.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name 2 | """Information registry to register operator information for compiler""" 3 | import tvm 4 | 5 | class OpPattern(object): 6 | """Operator generic patterns 7 | 8 | See Also 9 | -------- 10 | top.tag : Contains explanation of the tag type. 11 | """ 12 | # Elementwise operator 13 | ELEMWISE = 0 14 | # Broadcast operator 15 | BROADCAST = 1 16 | # Injective mapping 17 | INJECTIVE = 2 18 | # Comunication 19 | COMM_REDUCE = 3 20 | # Complex op, can still fuse ewise into it 21 | OUT_ELEMWISE_FUSABLE = 4 22 | # Not fusable opaque op 23 | OPAQUE = 8 24 | 25 | _register_compute = tvm.get_global_func("nnvm._register_compute") 26 | _register_schedule = tvm.get_global_func("nnvm._register_schedule") 27 | _register_pattern = tvm.get_global_func("nnvm._register_pattern") 28 | _register_alter_op_layout = tvm.get_global_func("nnvm.compiler._register_alter_op_layout") 29 | 30 | def register_compute(op_name, f=None, level=10): 31 | """Register compute function for operator 32 | 33 | Parameters 34 | ---------- 35 | op_name : str 36 | The name of operator 37 | 38 | f : function 39 | The schedule function 40 | 41 | level : int 42 | The priority level 43 | 44 | Returns 45 | ------- 46 | fregister : function 47 | Register function if f is not specified. 48 | """ 49 | def register(myf): 50 | """internal register function""" 51 | _register_compute(op_name, myf, level) 52 | return myf 53 | return register(f) if f else register 54 | 55 | 56 | def register_schedule(op_name, f=None, level=10): 57 | """Register schedule function for operator 58 | 59 | Parameters 60 | ---------- 61 | op_name : str 62 | The name of operator 63 | 64 | f : function 65 | The schedule function 66 | 67 | level : int 68 | The priority level 69 | 70 | Returns 71 | ------- 72 | fregister : function 73 | Register function if f is not specified. 74 | """ 75 | def register(myf): 76 | """internal register function""" 77 | _register_schedule(op_name, myf, level) 78 | return myf 79 | return register(f) if f else register 80 | 81 | 82 | def register_pattern(op_name, pattern, level=10): 83 | """Register pattern code for operator 84 | 85 | Parameters 86 | ---------- 87 | op_name : str 88 | The name of operator 89 | 90 | pattern : int 91 | The pattern code. 92 | 93 | level : int 94 | The priority level 95 | """ 96 | _register_pattern(op_name, pattern, level) 97 | 98 | 99 | def register_alter_op_layout(op_name, f=None, level=10): 100 | """Register alter layout function for operator 101 | 102 | Parameters 103 | ---------- 104 | op_name : str 105 | The name of operator 106 | 107 | f : function 108 | The schedule function 109 | 110 | level : int 111 | The priority level 112 | 113 | Returns 114 | ------- 115 | fregister : function 116 | Register function if f is not specified. 117 | """ 118 | def register(myf): 119 | """internal register function""" 120 | _register_alter_op_layout(op_name, myf, level) 121 | return myf 122 | return register(f) if f else register 123 | -------------------------------------------------------------------------------- /python/nnvm/top/transform.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name, unused-argument 2 | """Tensor transformation ops""" 3 | from __future__ import absolute_import 4 | 5 | import topi 6 | from .tensor import _fschedule_broadcast, _fschedule_injective 7 | from . import registry as reg 8 | from .registry import OpPattern 9 | 10 | # expand_dims 11 | reg.register_pattern("expand_dims", OpPattern.BROADCAST) 12 | reg.register_schedule("expand_dims", _fschedule_broadcast) 13 | 14 | # expand_like 15 | @reg.register_compute("expand_like") 16 | def compute_expand_like(attrs, inputs, _): 17 | """Compute definition of expand_like""" 18 | exclude = attrs.get_bool("exclude") 19 | axis = attrs.get_int_tuple("axis") 20 | if exclude: 21 | exclude_axis = (axis,) if isinstance(axis, int) else axis 22 | axis = [] 23 | for item in range(len(inputs[1].shape)): 24 | if item not in exclude_axis: 25 | axis.append(item) 26 | axis = tuple(axis) 27 | 28 | return topi.transform.expand_like(inputs[0], inputs[1], axis) 29 | reg.register_pattern("expand_like", OpPattern.BROADCAST) 30 | reg.register_schedule("expand_like", _fschedule_broadcast) 31 | 32 | # reshape_like 33 | @reg.register_compute("reshape_like") 34 | def compute_reshape_like(attrs, inputs, out_info): 35 | """Compute definition of reshape_like""" 36 | return topi.reshape(inputs[0], inputs[1].shape) 37 | reg.register_pattern("reshape_like", OpPattern.INJECTIVE) 38 | reg.register_schedule("reshape_like", _fschedule_injective) 39 | 40 | # transpose 41 | reg.register_pattern("transpose", OpPattern.INJECTIVE) 42 | reg.register_schedule("transpose", _fschedule_injective) 43 | 44 | # flip 45 | reg.register_pattern("flip", OpPattern.INJECTIVE) 46 | reg.register_schedule("flip", _fschedule_injective) 47 | 48 | # reshape 49 | reg.register_pattern("reshape", OpPattern.INJECTIVE) 50 | reg.register_schedule("reshape", _fschedule_injective) 51 | 52 | # squeeze 53 | reg.register_pattern("squeeze", OpPattern.INJECTIVE) 54 | reg.register_schedule("squeeze", _fschedule_injective) 55 | 56 | # concatenate 57 | reg.register_pattern("concatenate", OpPattern.INJECTIVE) 58 | reg.register_schedule("concatenate", _fschedule_injective) 59 | 60 | # split 61 | reg.register_pattern("split", OpPattern.INJECTIVE) 62 | reg.register_schedule("split", _fschedule_injective) 63 | -------------------------------------------------------------------------------- /python/nnvm/top/vision.py: -------------------------------------------------------------------------------- 1 | 2 | # pylint: disable=invalid-name, unused-argument 3 | """Definition of nn ops""" 4 | from __future__ import absolute_import 5 | 6 | import topi 7 | import tvm 8 | from . import registry as reg 9 | from .registry import OpPattern 10 | 11 | @reg.register_compute("yolo2_reorg") 12 | def compute_reorg(attrs, inputs, _): 13 | """Compute definition of reorg""" 14 | return topi.vision.reorg(inputs[0], attrs.get_int("stride")) 15 | 16 | @reg.register_schedule("yolo2_reorg") 17 | def schedule_reorg(attrs, outs, target): 18 | """Schedule definition of reorg""" 19 | with tvm.target.create(target): 20 | return topi.generic.schedule_injective(outs) 21 | 22 | reg.register_pattern("yolo2_reorg", OpPattern.INJECTIVE) 23 | 24 | @reg.register_compute("yolo2_region") 25 | def compute_region(attrs, inputs, _): 26 | """Compute definition of region""" 27 | n = attrs.get_int("n") 28 | classes = attrs.get_int("classes") 29 | coords = attrs.get_int("coords") 30 | background = attrs.get_int("background") 31 | softmax = attrs.get_int("softmax") 32 | return topi.vision.yolo2.region(inputs[0], n, classes, coords, background, softmax) 33 | 34 | @reg.register_schedule("yolo2_region") 35 | def schedule_region(attrs, outs, target): 36 | """Schedule definition of region""" 37 | with tvm.target.create(target): 38 | return topi.generic.vision.schedule_region(outs) 39 | 40 | reg.register_pattern("yolo2_region", OpPattern.OPAQUE) 41 | -------------------------------------------------------------------------------- /python/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from setuptools import find_packages 4 | from distutils.core import setup 5 | 6 | def config_cython(): 7 | # temporary disable cython for now 8 | # as NNVM uses local DLL build 9 | return [] 10 | try: 11 | from Cython.Build import cythonize 12 | from distutils.extension import Extension 13 | if sys.version_info >= (3, 0): 14 | subdir = "_cy3" 15 | else: 16 | subdir = "_cy2" 17 | ret = [] 18 | path = "nnvm/cython" 19 | 20 | for fn in os.listdir(path): 21 | if not fn.endswith(".pyx"): 22 | continue 23 | ret.append(Extension( 24 | "nnvm/%s/%s" % (subdir, fn[:-4]), 25 | ["nnvm/cython/%s" % fn], 26 | include_dirs=["../include/"], 27 | language="c++")) 28 | return cythonize(ret) 29 | except: 30 | print("Cython is not installed, will compile without cython module") 31 | return [] 32 | 33 | # We can not import `libinfo.py` in setup.py directly since __init__.py 34 | # Will be invoked which introduces dependences 35 | CURRENT_DIR = os.path.dirname(__file__) 36 | libinfo_py = os.path.join(CURRENT_DIR, './nnvm/libinfo.py') 37 | libinfo = {'__file__': libinfo_py} 38 | exec(compile(open(libinfo_py, "rb").read(), libinfo_py, 'exec'), libinfo, libinfo) 39 | 40 | LIB_PATH = libinfo['find_lib_path']() 41 | _, LIB_NAME = os.path.split(LIB_PATH[0]) 42 | __version__ = libinfo['__version__'] 43 | curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) 44 | rpath = os.path.relpath(LIB_PATH[0], curr_path) 45 | 46 | setup(name='nnvm', 47 | version=__version__, 48 | description="NNVM: Open Compiler for AI Frameworks", 49 | zip_safe=False, 50 | install_requires=[ 51 | 'numpy' 52 | ], 53 | packages=find_packages(), 54 | url='https://github.com/dmlc/nnvm', 55 | include_package_data=True, 56 | data_files=[('nnvm', [rpath])]) 57 | 58 | -------------------------------------------------------------------------------- /src/README.md: -------------------------------------------------------------------------------- 1 | Project Structure 2 | ================= 3 | 4 | The following components are operator invariant. 5 | 6 | - c_api: NNVM C API 7 | - core: NNVM core data structure 8 | - pass: NNVM pass 9 | 10 | The following components are generic NNVM compiler and defines tensor operator set 11 | 12 | - top: NNVM core tensor operators 13 | - compiler: NNVM compiler toolchain 14 | -------------------------------------------------------------------------------- /src/c_api/c_api_common.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file c_api_error.h 4 | * \brief Common fields of all C APIs 5 | */ 6 | #ifndef NNVM_C_API_C_API_COMMON_H_ 7 | #define NNVM_C_API_C_API_COMMON_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | /*! \brief macro to guard beginning and end section of all functions */ 18 | #define API_BEGIN() try { 19 | /*! \brief every function starts with API_BEGIN(); 20 | and finishes with API_END() or API_END_HANDLE_ERROR */ 21 | #define API_END() } catch(dmlc::Error &_except_) { return NNAPIHandleException(_except_); } return 0; // NOLINT(*) 22 | /*! 23 | * \brief every function starts with API_BEGIN(); 24 | * and finishes with API_END() or API_END_HANDLE_ERROR 25 | * The finally clause contains procedure to cleanup states when an error happens. 26 | */ 27 | #define API_END_HANDLE_ERROR(Finalize) } catch(dmlc::Error &_except_) { Finalize; return NNAPIHandleException(_except_); } return 0; // NOLINT(*) 28 | 29 | 30 | /*! \brief entry to to easily hold returning information */ 31 | struct NNAPIThreadLocalEntry { 32 | /*! \brief result holder for returning string */ 33 | std::string ret_str; 34 | /*! \brief result holder for returning strings */ 35 | std::vector ret_vec_str; 36 | /*! \brief result holder for returning string pointers */ 37 | std::vector ret_vec_charp; 38 | /*! \brief result holder for returning handles */ 39 | std::vector ret_handles; 40 | /*! \brief argument holder to hold symbol */ 41 | std::unordered_map kwarg_symbol; 42 | }; 43 | 44 | /*! \brief Thread local store that can be used to hold return values. */ 45 | typedef dmlc::ThreadLocalStore NNAPIThreadLocalStore; 46 | 47 | /*! 48 | * \brief handle exception throwed out 49 | * \param e the exception 50 | * \return the return value of API after exception is handled 51 | */ 52 | inline int NNAPIHandleException(const dmlc::Error &e) { 53 | NNAPISetLastError(e.what()); 54 | return -1; 55 | } 56 | 57 | #endif // NNVM_C_API_C_API_COMMON_H_ 58 | -------------------------------------------------------------------------------- /src/c_api/c_api_error.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file c_api_error.cc 4 | * \brief C error handling 5 | */ 6 | #include 7 | #include "./c_api_common.h" 8 | 9 | struct ErrorEntry { 10 | std::string last_error; 11 | }; 12 | 13 | typedef dmlc::ThreadLocalStore NNAPIErrorStore; 14 | 15 | const char *NNGetLastError() { 16 | return NNAPIErrorStore::Get()->last_error.c_str(); 17 | } 18 | 19 | void NNAPISetLastError(const char* msg) { 20 | NNAPIErrorStore::Get()->last_error = msg; 21 | } 22 | -------------------------------------------------------------------------------- /src/c_api/c_api_graph.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file c_api_graph.cc 4 | * \brief C API related to Graph IR. 5 | */ 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include "./c_api_common.h" 13 | 14 | using namespace nnvm; 15 | 16 | int NNGraphCreate(SymbolHandle symbol, GraphHandle *graph) { 17 | Graph* g = new Graph(); 18 | API_BEGIN(); 19 | g->outputs = static_cast(symbol)->outputs; 20 | *graph = g; 21 | API_END_HANDLE_ERROR(delete g); 22 | } 23 | 24 | int NNGraphFree(GraphHandle handle) { 25 | API_BEGIN(); 26 | delete static_cast(handle); 27 | API_END(); 28 | } 29 | 30 | int NNGraphGetSymbol(GraphHandle graph, SymbolHandle *symbol) { 31 | Symbol* s = new Symbol(); 32 | API_BEGIN(); 33 | s->outputs = static_cast(graph)->outputs; 34 | *symbol = s; 35 | API_END_HANDLE_ERROR(delete s); 36 | } 37 | 38 | int NNGraphSetNodeEntryListAttr_(GraphHandle handle, 39 | const char* key, 40 | SymbolHandle list) { 41 | API_BEGIN(); 42 | Symbol* s = static_cast(list); 43 | Graph* g = static_cast(handle); 44 | g->attrs[std::string(key)] 45 | = std::make_shared(s->outputs); 46 | API_END(); 47 | } 48 | 49 | int NNGraphSetJSONAttr(GraphHandle handle, 50 | const char* key, 51 | const char* json_value) { 52 | API_BEGIN(); 53 | Graph* g = static_cast(handle); 54 | std::string temp(json_value); 55 | std::istringstream is(temp); 56 | dmlc::JSONReader reader(&is); 57 | nnvm::any value; 58 | reader.Read(&value); 59 | g->attrs[std::string(key)] = std::make_shared(std::move(value)); 60 | API_END(); 61 | } 62 | 63 | int NNGraphGetJSONAttr(GraphHandle handle, 64 | const char* key, 65 | const char** json_out, 66 | int *success) { 67 | NNAPIThreadLocalEntry *ret = NNAPIThreadLocalStore::Get(); 68 | API_BEGIN(); 69 | Graph* g = static_cast(handle); 70 | std::string skey(key); 71 | auto it = g->attrs.find(skey); 72 | if (it != g->attrs.end()) { 73 | std::ostringstream os; 74 | dmlc::JSONWriter writer(&os); 75 | writer.Write(*it->second.get()); 76 | ret->ret_str = os.str(); 77 | *json_out = (ret->ret_str).c_str(); 78 | *success = 1; 79 | } else { 80 | *success = 0; 81 | } 82 | API_END(); 83 | } 84 | 85 | int NNGraphApplyPasses(GraphHandle src, 86 | nn_uint num_pass, 87 | const char** pass_names, 88 | GraphHandle *dst) { 89 | Graph* g = new Graph(); 90 | API_BEGIN(); 91 | std::vector vpass; 92 | for (nn_uint i = 0; i < num_pass; ++i) { 93 | vpass.emplace_back(std::string(pass_names[i])); 94 | } 95 | *g = ApplyPasses(*static_cast(src), vpass); 96 | *dst = g; 97 | API_END_HANDLE_ERROR(delete g); 98 | } 99 | -------------------------------------------------------------------------------- /src/compiler/compile_engine.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file compile_engine.h 4 | * \brief Internal engine to compile a subgraph fragment and cache compilation. 5 | */ 6 | #ifndef NNVM_COMPILER_COMPILE_ENGINE_H_ 7 | #define NNVM_COMPILER_COMPILE_ENGINE_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include "./graph_hash.h" 22 | 23 | namespace nnvm { 24 | namespace compiler { 25 | 26 | /*! \brief A TVM Node to represent compiled graph function */ 27 | struct GraphFuncNode : public tvm::Node { 28 | /* \brief compiled target */ 29 | std::string target; 30 | /*! \brief Function name */ 31 | std::string func_name; 32 | /* \brief The inputs to the function */ 33 | tvm::Array inputs; 34 | /* \brief The outputs to the function */ 35 | tvm::Array outputs; 36 | /*! \brief The lowered functions */ 37 | tvm::Array funcs; 38 | 39 | void VisitAttrs(tvm::AttrVisitor* v) final { 40 | v->Visit("target", &target); 41 | v->Visit("func_name", &func_name); 42 | v->Visit("inputs", &inputs); 43 | v->Visit("outputs", &outputs); 44 | v->Visit("funcs", &funcs); 45 | } 46 | 47 | static constexpr const char* _type_key = "GraphFunc"; 48 | TVM_DECLARE_NODE_TYPE_INFO(GraphFuncNode, tvm::Node); 49 | }; 50 | 51 | TVM_DEFINE_NODE_REF(GraphFunc, GraphFuncNode); 52 | 53 | /*! \brief Cache Entry in the graph */ 54 | struct GraphCacheEntryNode : public tvm::Node { 55 | /*! \brief The graph function */ 56 | GraphFunc graph_func; 57 | /*! \brief Usage statistics */ 58 | int use_count{0}; 59 | /*! \brief Index of the master node for calling schedule*/ 60 | int master_idx; 61 | 62 | void VisitAttrs(tvm::AttrVisitor* v) final { 63 | v->Visit("graph_func", &graph_func); 64 | v->Visit("use_count", &use_count); 65 | v->Visit("master_idx", &master_idx); 66 | } 67 | static constexpr const char* _type_key = "GraphCacheEntry"; 68 | TVM_DECLARE_NODE_TYPE_INFO(GraphCacheEntryNode, tvm::Node); 69 | }; 70 | 71 | class GraphCacheEntry : public ::tvm::NodeRef { 72 | public: 73 | GraphCacheEntry() {} 74 | explicit GraphCacheEntry(std::shared_ptr<::tvm::Node> n) : NodeRef(n) {} 75 | GraphCacheEntryNode* operator->() { 76 | return static_cast(node_.get()); 77 | } 78 | using ContainerType = GraphCacheEntryNode; 79 | }; 80 | 81 | /*! 82 | * \brief Call compile engine to lower a graph with given inputs. 83 | * 84 | * \param graph The graph to be compiled 85 | * \param inputs The input specification. 86 | * \param target The build target 87 | * \param master_idx The index of master node for calling schedule 88 | * 89 | * \return func A lowered tvm function. 90 | */ 91 | GraphFunc GraphLower(Graph graph, 92 | const Array& inputs, 93 | const std::string& target, 94 | int master_idx); 95 | 96 | /*! 97 | * \brief Get type flag from TVM Type 98 | * 99 | * \param type the tvm type 100 | * \return corresponding DLDataType 101 | */ 102 | int GetTypeFlag(tvm::Type type); 103 | 104 | /*! 105 | * \brief Get TVM Type from type flag 106 | * 107 | * \param type_flag the type flag 108 | * \return corresponding TVM type 109 | */ 110 | tvm::Type GetTVMType(int type_flag); 111 | 112 | } // namespace compiler 113 | } // namespace nnvm 114 | 115 | #endif // NNVM_COMPILER_COMPILE_ENGINE_H_ 116 | -------------------------------------------------------------------------------- /src/compiler/graph_hash.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file graph_hash.h 4 | * \brief The graph hashing function. 5 | */ 6 | #ifndef NNVM_COMPILER_GRAPH_HASH_H_ 7 | #define NNVM_COMPILER_GRAPH_HASH_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace nnvm { 15 | namespace compiler { 16 | 17 | class GraphKey; 18 | 19 | /*! \brief Key to a graph compiler cache */ 20 | struct GraphKeyNode : public tvm::Node { 21 | /*! \brief The graph structure */ 22 | Graph graph; 23 | /* \brief The inputs to the function */ 24 | tvm::Array inputs; 25 | /*! \brief The target */ 26 | std::string target; 27 | // Cached internal hash key, invisible to the user. 28 | // The graph hash key is ensured always not to be 0 29 | mutable size_t cache_hash_key_{0}; 30 | 31 | void VisitAttrs(tvm::AttrVisitor* v) final { 32 | v->Visit("inputs", &inputs); 33 | v->Visit("target", &target); 34 | } 35 | 36 | static GraphKey make(Graph graph, 37 | tvm::Array inputs, 38 | std::string target); 39 | static constexpr const char* _type_key = "GraphKey"; 40 | TVM_DECLARE_NODE_TYPE_INFO(GraphKeyNode, tvm::Node); 41 | }; 42 | 43 | TVM_DEFINE_NODE_REF(GraphKey, GraphKeyNode); 44 | 45 | /*! \brief Hashing function for graph key */ 46 | struct GraphKeyHash { 47 | size_t operator()(const GraphKey& gkey) const { 48 | return Hash(gkey); 49 | } 50 | static size_t Hash(const GraphKey& gkey); 51 | }; 52 | 53 | /*! \brief function for graph key */ 54 | struct GraphKeyEqual { 55 | bool operator()(const GraphKey& a, 56 | const GraphKey& b) const { 57 | return Equal(a, b); 58 | } 59 | static bool Equal(const GraphKey& a, const GraphKey& b); 60 | }; 61 | 62 | /*! 63 | * \brief Create a hash code for a given graph. 64 | * \return The hash code of the graph. 65 | */ 66 | size_t GraphHash(const Graph& graph); 67 | 68 | /*! 69 | * \brief Compare two graphs 70 | * return empty string if they are equal 71 | * otherwise return error message 72 | * \param a The first graph. 73 | * \param b The second graph. 74 | * \return empty string if they are equal, otherwise return error message. 75 | */ 76 | std::string GraphDeepCompare(const Graph& a, 77 | const Graph& b, 78 | bool compare_variable_attr); 79 | } // namespace compiler 80 | } // namespace nnvm 81 | 82 | #endif // NNVM_COMPILER_GRAPH_HASH_H_ 83 | -------------------------------------------------------------------------------- /src/compiler/graph_runtime.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file graph_runtime.h 4 | * \brief Interface code with TVM graph runtime. 5 | */ 6 | #ifndef NNVM_COMPILER_GRAPH_RUNTIME_H_ 7 | #define NNVM_COMPILER_GRAPH_RUNTIME_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | namespace nnvm { 14 | namespace compiler { 15 | 16 | /*! \brief Magic number for NDArray file */ 17 | constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F; 18 | /*! \brief Magic number for NDArray list file */ 19 | constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; 20 | 21 | struct TVMOpParam : public dmlc::Parameter { 22 | std::string func_name; 23 | uint32_t num_inputs; 24 | uint32_t num_outputs; 25 | uint32_t flatten_data; 26 | 27 | DMLC_DECLARE_PARAMETER(TVMOpParam) { 28 | DMLC_DECLARE_FIELD(func_name); 29 | DMLC_DECLARE_FIELD(num_inputs).set_default(1); 30 | DMLC_DECLARE_FIELD(num_outputs).set_default(1); 31 | DMLC_DECLARE_FIELD(flatten_data).set_default(0); 32 | } 33 | }; 34 | 35 | } // namespace compiler 36 | } // namespace nnvm 37 | #endif // NNVM_COMPILER_GRAPH_RUNTIME_H_ 38 | -------------------------------------------------------------------------------- /src/compiler/node_attr.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file node_attr.h 4 | * \brief utility to access node attributes 5 | */ 6 | #ifndef NNVM_COMPILER_NODE_ATTR_H_ 7 | #define NNVM_COMPILER_NODE_ATTR_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace nnvm { 15 | namespace compiler { 16 | 17 | using AttrDict = std::unordered_map; 18 | /*! 19 | * \brief Get canonicalized attr dict from node 20 | * \param attrs The node attrs 21 | * \return The attribute dict 22 | */ 23 | inline AttrDict GetAttrDict(const NodeAttrs& attrs) { 24 | static auto& fgetdict = nnvm::Op::GetAttr("FGetAttrDict"); 25 | if (fgetdict.count(attrs.op)) { 26 | return fgetdict[attrs.op](attrs); 27 | } else { 28 | return attrs.dict; 29 | } 30 | } 31 | 32 | } // namespace compiler 33 | } // namespace nnvm 34 | #endif // NNVM_COMPILER_NODE_ATTR_H_ 35 | -------------------------------------------------------------------------------- /src/compiler/pattern_util.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file pattern_util.h 4 | * \brief Utilities for doing various pattern matching in graph. 5 | */ 6 | #ifndef NNVM_COMPILER_PATTERN_UTIL_H_ 7 | #define NNVM_COMPILER_PATTERN_UTIL_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace nnvm { 15 | namespace compiler { 16 | 17 | /*! 18 | * \brief find axis in oshape, such that: 19 | * bias_shape = [1,1, ... oshape[axis], 1,1,] 20 | * 21 | * This is used to detect bias or scaling factor on channel dimension. 22 | * \param oshape The output shape 23 | * \param bias_shape The shape of bias or scaling factor. 24 | * \return Pair of matched axis in o shape and bias_shape if found. 25 | */ 26 | inline std::pair MatchBroadcast1DAxis( 27 | const TShape& oshape, const TShape& bias_shape) { 28 | dim_t axis_dim = bias_shape.ndim(); 29 | for (dim_t i = bias_shape.ndim(); i != 0; --i, --axis_dim) { 30 | if (bias_shape[i - 1] != 1) break; 31 | } 32 | // everything is 1 33 | if (axis_dim == 0) { 34 | return {oshape.ndim() - bias_shape.ndim(), 0}; 35 | } 36 | axis_dim = axis_dim - 1; 37 | // The bias shape is not 1D 38 | for (dim_t i = 0; i < axis_dim; ++i) { 39 | if (bias_shape[i] != 1) return {-1, -1}; 40 | } 41 | int axis = static_cast( 42 | oshape.ndim() - bias_shape.ndim() + axis_dim); 43 | if (oshape[axis] != bias_shape[axis_dim]) return {-1, -1}; 44 | return {axis, axis_dim}; 45 | } 46 | 47 | /*! 48 | * \brief Expand bias dimension to match needed axis. 49 | * 50 | * \param bias The bias NodeEntry 51 | * \param out_dim output dimension. 52 | * \param bias_dim The current bias dimension. 53 | * \param axis The axis we want to match on. 54 | */ 55 | inline NodeEntry 56 | ExpandBiasToMatchAxis(NodeEntry bias, 57 | int out_dim, 58 | int bias_dim, 59 | int axis) { 60 | if (bias_dim != 1) { 61 | bias = MakeNode("squeeze", bias.node->attrs.name + "_sqz", {bias}); 62 | } 63 | int num_pad_axis = out_dim - axis - 1; 64 | if (num_pad_axis > 0) { 65 | std::unordered_map kwargs{ 66 | {"axis", "1"}, 67 | {"num_newaxis", std::to_string(num_pad_axis)}}; 68 | return MakeNode("expand_dims", bias.node->attrs.name + "_expand", 69 | {bias}, kwargs); 70 | 71 | } else { 72 | return bias; 73 | } 74 | } 75 | 76 | /*! 77 | * \brief Get the reference count of each node. 78 | * \param idx The IndexedGraph 79 | * \return ref_count vector of length number nodes. 80 | */ 81 | inline std::vector 82 | GetNodeRefCounts(const IndexedGraph& idx) { 83 | std::vector ref_count(idx.num_nodes(), 0); 84 | for (uint32_t nid = 0; nid < idx.num_nodes(); ++nid) { 85 | const auto& inode = idx[nid]; 86 | if (inode.source->is_variable()) continue; 87 | for (const auto& e : inode.inputs) { 88 | ++ref_count[e.node_id]; 89 | } 90 | } 91 | for (const auto& e : idx.outputs()) { 92 | // this line will realize all the outputs 93 | ref_count[e.node_id] += 1; 94 | } 95 | return ref_count; 96 | } 97 | } // namespace compiler 98 | } // namespace nnvm 99 | #endif // NNVM_COMPILER_PATTERN_UTIL_H_ 100 | -------------------------------------------------------------------------------- /src/compiler/precompute_prune.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file precompute_prune.cc 4 | * \brief Split the graph into a pre-compute graph and a execution graph. 5 | * 6 | * The pre-compute graph outputs parameters that can be taken 7 | * by execution graph during execution phase. 8 | */ 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace nnvm { 17 | namespace compiler { 18 | 19 | nnvm::Graph PrecomputePrune(nnvm::Graph src) { 20 | const auto& plist 21 | = src.GetAttr >("param_name_list"); 22 | std::unordered_set params(plist.begin(), plist.end()); 23 | 24 | std::unordered_set pruned; 25 | nnvm::NodeEntryMap entry_var; 26 | std::unordered_set unique_name; 27 | // number of edges that are not variable 28 | int non_var_edge = 0; 29 | 30 | auto replace_pruned_entry = [&] (const NodeEntry& e) { 31 | if (!entry_var.count(e)) { 32 | if (!e.node->is_variable()) { 33 | ++non_var_edge; 34 | } 35 | nnvm::NodePtr var = nnvm::Node::Create(); 36 | var->attrs.name = e.node->attrs.name; 37 | if (e.version) { 38 | var->attrs.name += "_" + std::to_string(e.version); 39 | } 40 | if (e.node->num_outputs() != 1) { 41 | var->attrs.name += "_output" + std::to_string(e.index); 42 | } 43 | entry_var.emplace(e, var); 44 | CHECK(!unique_name.count(var->attrs.name)); 45 | unique_name.insert(var->attrs.name); 46 | return nnvm::NodeEntry{var, 0, 0}; 47 | } else { 48 | return nnvm::NodeEntry{entry_var.at(e), 0, 0}; 49 | } 50 | }; 51 | 52 | DFSVisit(src.outputs, [&](const nnvm::NodePtr& n) { 53 | bool can_be_pruned = true; 54 | if (n->is_variable()) { 55 | if (params.count(n->attrs.name)) { 56 | pruned.emplace(n.get()); 57 | } 58 | can_be_pruned = false; 59 | } 60 | 61 | for (const auto& e : n->inputs) { 62 | if (!pruned.count(e.node.get())) { 63 | can_be_pruned = false; 64 | } 65 | } 66 | if (can_be_pruned) { 67 | pruned.emplace(n.get()); 68 | } else { 69 | // scan again to find edge nodes, skip variables 70 | for (auto& e : n->inputs) { 71 | if (pruned.count(e.node.get())) { 72 | e = replace_pruned_entry(e); 73 | } 74 | } 75 | } 76 | }); 77 | 78 | // nothing being pruned. 79 | if (non_var_edge == 0) { 80 | return src; 81 | } 82 | 83 | for (auto& e : src.outputs) { 84 | if (pruned.count(e.node.get())) { 85 | e = replace_pruned_entry(e); 86 | } 87 | } 88 | 89 | nnvm::Graph pre_graph; 90 | pre_graph.outputs.reserve(entry_var.size()); 91 | std::vector output_names; 92 | output_names.reserve(entry_var.size()); 93 | 94 | for (auto kv : entry_var) { 95 | pre_graph.outputs.emplace_back(kv.first); 96 | output_names.emplace_back(kv.second->attrs.name); 97 | } 98 | // new parameter list 99 | pre_graph.attrs["output_names"] = 100 | std::make_shared(std::move(output_names)); 101 | src.attrs["precompute_graph"] = 102 | std::make_shared(std::move(pre_graph)); 103 | return src; 104 | } 105 | 106 | NNVM_REGISTER_PASS(PrecomputePrune) 107 | .set_body(PrecomputePrune); 108 | } // namespace compiler 109 | } // namespace nnvm 110 | -------------------------------------------------------------------------------- /src/compiler/simplify_inference.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file simplify_inference.cc 4 | * \author Ziheng Jiang 5 | */ 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include "./graph_transform.h" 13 | #include "./pattern_util.h" 14 | 15 | namespace nnvm { 16 | namespace compiler { 17 | 18 | std::vector 19 | BatchNormToInferUnpack(const nnvm::NodeAttrs& attrs, 20 | nnvm::NodeEntry data, 21 | nnvm::NodeEntry gamma, 22 | nnvm::NodeEntry beta, 23 | nnvm::NodeEntry moving_mean, 24 | nnvm::NodeEntry moving_var, 25 | TShape dshape, 26 | TShape bshape) { 27 | CHECK_NE(dshape.ndim(), 0); 28 | CHECK(attrs.op); 29 | static const Op* bn_op = Op::Get("batch_norm"); 30 | CHECK(attrs.op == bn_op); 31 | const auto& param = nnvm::get(attrs.parsed); 32 | std::string bn_name = attrs.name; 33 | 34 | // transform batch_norm(data) to scale * data + shift 35 | NodeEntry var_add_eps = MakeNode( 36 | "__add_scalar__", bn_name + "_add_eps", 37 | {moving_var}, {{"scalar", std::to_string(param.epsilon)}}); 38 | 39 | NodeEntry sqrt = MakeNode( 40 | "sqrt", bn_name + "_sqrt", {var_add_eps}); 41 | 42 | NodeEntry scale = MakeNode( 43 | "__rdiv_scalar__", bn_name + "_div", 44 | {sqrt}, {{"scalar", "1"}}); 45 | 46 | if (param.scale) { 47 | scale = MakeNode( 48 | "elemwise_mul", bn_name + "_gamma_mul_div", 49 | {scale, gamma}); 50 | } 51 | 52 | NodeEntry neg_mean = MakeNode( 53 | "negative", bn_name + "_neg_mean", {moving_mean}); 54 | 55 | NodeEntry shift = MakeNode( 56 | "elemwise_mul", bn_name + "_neg_mean_mul_a", 57 | {neg_mean, scale}); 58 | 59 | if (param.center) { 60 | shift = MakeNode( 61 | "elemwise_add", bn_name + "_add_beta", {shift, beta}); 62 | } 63 | int axis = param.axis; 64 | scale = ExpandBiasToMatchAxis(scale, dshape.ndim()-bshape.ndim()+1, 1, axis); 65 | shift = ExpandBiasToMatchAxis(shift, dshape.ndim()-bshape.ndim()+1, 1, axis); 66 | 67 | NodeEntry out = MakeNode("broadcast_mul", bn_name + "_a_mul_data", 68 | {data, scale}); 69 | out = MakeNode("broadcast_add", bn_name + "_out", 70 | {out, shift}); 71 | // It is invalid to ref the other values of BN after inference transform. 72 | NodeEntry undef = MakeNode("__undef__", "undef", {}); 73 | return {out, undef, undef}; 74 | } 75 | 76 | Graph SimplifyInference(nnvm::Graph src) { 77 | // Get attributes from the graph 78 | const IndexedGraph& idx = src.indexed_graph(); 79 | const ShapeVector& shape_vec = src.GetAttr("shape"); 80 | auto transform = [&](uint32_t nid, const NodePtr& n, std::vector* ret) { 81 | if (n->is_variable()) return false; 82 | static const Op* bn_op = Op::Get("batch_norm"); 83 | static const Op* dropout_op = Op::Get("dropout"); 84 | if (n->op() == bn_op) { 85 | *ret = BatchNormToInferUnpack( 86 | n->attrs, 87 | n->inputs[0], 88 | n->inputs[1], 89 | n->inputs[2], 90 | n->inputs[3], 91 | n->inputs[4], 92 | shape_vec[idx.entry_id(nid, 0)], 93 | shape_vec[idx.entry_id(nid, 1)]); 94 | return true; 95 | } else if (n->op() == dropout_op) { 96 | NodeEntry undef = MakeNode("__undef__", "undef", {}); 97 | *ret = {n->inputs[0], undef}; 98 | return true; 99 | } else { 100 | return false; 101 | } 102 | }; 103 | return GraphTransform(src, transform); 104 | } 105 | 106 | NNVM_REGISTER_PASS(SimplifyInference) 107 | .set_body(SimplifyInference) 108 | .set_change_graph(true); 109 | 110 | } // namespace compiler 111 | } // namespace nnvm 112 | -------------------------------------------------------------------------------- /src/core/graph.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file graph_attr_types.cc 4 | * \brief Graph node data structure. 5 | */ 6 | #include 7 | #include 8 | #include 9 | 10 | namespace nnvm { 11 | 12 | const IndexedGraph& Graph::indexed_graph() const { 13 | if (indexed_graph_ == nullptr) { 14 | indexed_graph_.reset(new IndexedGraph(*this)); 15 | } 16 | return *indexed_graph_; 17 | } 18 | 19 | // implement constructor from graph 20 | IndexedGraph::IndexedGraph(const Graph &g) { 21 | entry_rptr_.push_back(0); 22 | std::vector inputs_rptr{0}, control_rptr{0}; 23 | 24 | DFSVisit(g.outputs, [this, &inputs_rptr, &control_rptr] 25 | (const NodePtr& n) { 26 | CHECK_LT(nodes_.size(), std::numeric_limits::max()); 27 | uint32_t nid = static_cast(nodes_.size()); 28 | // nodes_ 29 | IndexedGraph::Node new_node; 30 | new_node.source = n.get(); 31 | new_node.weak_ref = n; 32 | nodes_.emplace_back(std::move(new_node)); 33 | // arg_nodes_ 34 | if (n->is_variable()) { 35 | input_nodes_.push_back(nid); 36 | } 37 | // node2index_ 38 | node2index_[n.get()] = nid; 39 | // entry rptr 40 | entry_rptr_.push_back(entry_rptr_.back() + n->num_outputs()); 41 | // input entries 42 | for (const auto& e : n->inputs) { 43 | auto it = node2index_.find(e.node.get()); 44 | CHECK(it != node2index_.end() && it->first == e.node.get()); 45 | input_entries_.emplace_back(NodeEntry{it->second, e.index, e.version}); 46 | } 47 | inputs_rptr.push_back(input_entries_.size()); 48 | // control deps 49 | for (const auto& nptr : n->control_deps) { 50 | auto it = node2index_.find(nptr.get()); 51 | CHECK(it != node2index_.end() && it->first == nptr.get()); 52 | control_deps_.push_back(it->second); 53 | } 54 | control_rptr.push_back(control_deps_.size()); 55 | }); 56 | 57 | for (const auto& e : g.outputs) { 58 | outputs_.emplace_back(NodeEntry{ 59 | node2index_.at(e.node.get()), e.index, e.version}); 60 | } 61 | 62 | static auto& fmutate_inputs = Op::GetAttr("FMutateInputs"); 63 | // setup array view 64 | // input_entries_ and control_rptr must not change after this step. 65 | const NodeEntry* iptr = dmlc::BeginPtr(input_entries_); 66 | for (size_t nid = 0; nid < nodes_.size(); ++nid) { 67 | nodes_[nid].inputs = array_view( 68 | iptr + inputs_rptr[nid], iptr + inputs_rptr[nid + 1]); 69 | if (nodes_[nid].source->op() != nullptr && 70 | fmutate_inputs.count(nodes_[nid].source->op())) { 71 | for (uint32_t i : fmutate_inputs[nodes_[nid].source->op()](nodes_[nid].source->attrs)) { 72 | mutable_input_nodes_.insert(nodes_[nid].inputs[i].node_id); 73 | } 74 | } 75 | } 76 | const uint32_t* cptr = dmlc::BeginPtr(control_deps_); 77 | for (size_t nid = 0; nid < nodes_.size(); ++nid) { 78 | nodes_[nid].control_deps = array_view( 79 | cptr + control_rptr[nid], cptr + control_rptr[nid + 1]); 80 | } 81 | } 82 | 83 | } // namespace nnvm 84 | -------------------------------------------------------------------------------- /src/core/node.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file node.cc 4 | * \brief Graph node data structure. 5 | */ 6 | #include 7 | 8 | namespace nnvm { 9 | 10 | Node::~Node() { 11 | if (inputs.size() != 0) { 12 | // explicit deletion via DFS 13 | // this is used to avoid stackoverflow caused by chain of deletions 14 | std::vector stack{this}; 15 | std::vector to_delete; 16 | while (!stack.empty()) { 17 | Node* n = stack.back(); 18 | stack.pop_back(); 19 | for (NodeEntry& e : n->inputs) { 20 | if (e.node.unique()) { 21 | stack.push_back(e.node.get()); 22 | to_delete.emplace_back(std::move(e.node)); 23 | } else { 24 | e.node.reset(); 25 | } 26 | } 27 | for (NodePtr& sp : n->control_deps) { 28 | if (sp.unique()) { 29 | stack.push_back(sp.get()); 30 | to_delete.emplace_back(std::move(sp)); 31 | } else { 32 | sp.reset(); 33 | } 34 | } 35 | n->inputs.clear(); 36 | } 37 | } 38 | } 39 | 40 | NodePtr Node::Create() { 41 | return std::make_shared(); 42 | } 43 | 44 | } // namespace nnvm 45 | -------------------------------------------------------------------------------- /src/core/op.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file op.cc 4 | * \brief Support for operator registry. 5 | */ 6 | #include 7 | #include 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace dmlc { 15 | // enable registry 16 | DMLC_REGISTRY_ENABLE(nnvm::Op); 17 | } // namespace dmlc 18 | 19 | namespace nnvm { 20 | 21 | // single manager of operator information. 22 | struct OpManager { 23 | // mutex to avoid registration from multiple threads. 24 | // recursive is needed for trigger(which calls UpdateAttrMap) 25 | std::recursive_mutex mutex; 26 | // global operator counter 27 | std::atomic op_counter{0}; 28 | // storage of additional attribute table. 29 | std::unordered_map > attr; 30 | // storage of existing triggers 31 | std::unordered_map > > tmap; 32 | // group of each operator. 33 | std::vector > op_group; 34 | // get singleton of the 35 | static OpManager* Global() { 36 | static OpManager inst; 37 | return &inst; 38 | } 39 | }; 40 | 41 | // constructor 42 | Op::Op() { 43 | OpManager* mgr = OpManager::Global(); 44 | index_ = mgr->op_counter++; 45 | } 46 | 47 | Op& Op::add_alias(const std::string& alias) { // NOLINT(*) 48 | dmlc::Registry::Get()->AddAlias(this->name, alias); 49 | return *this; 50 | } 51 | 52 | // find operator by name 53 | const Op* Op::Get(const std::string& name) { 54 | const Op* op = dmlc::Registry::Find(name); 55 | CHECK(op != nullptr) 56 | << "Operator " << name << " is not registered"; 57 | return op; 58 | } 59 | 60 | // Get attribute map by key 61 | const any* Op::GetAttrMap(const std::string& key) { 62 | auto& dict = OpManager::Global()->attr; 63 | auto it = dict.find(key); 64 | if (it != dict.end()) { 65 | return it->second.get(); 66 | } else { 67 | return nullptr; 68 | } 69 | } 70 | 71 | // update attribute map 72 | void Op::UpdateAttrMap(const std::string& key, 73 | std::function updater) { 74 | OpManager* mgr = OpManager::Global(); 75 | std::lock_guard(mgr->mutex); 76 | std::unique_ptr& value = mgr->attr[key]; 77 | if (value.get() == nullptr) value.reset(new any()); 78 | if (updater != nullptr) updater(value.get()); 79 | } 80 | 81 | void Op::AddGroupTrigger(const std::string& group_name, 82 | std::function trigger) { 83 | OpManager* mgr = OpManager::Global(); 84 | std::lock_guard(mgr->mutex); 85 | auto& tvec = mgr->tmap[group_name]; 86 | tvec.push_back(trigger); 87 | auto& op_group = mgr->op_group; 88 | for (const Op* op : dmlc::Registry::List()) { 89 | if (op->index_ < op_group.size() && 90 | op_group[op->index_].count(group_name) != 0) { 91 | trigger((Op*)op); // NOLINT(*) 92 | } 93 | } 94 | } 95 | 96 | Op& Op::include(const std::string& group_name) { 97 | OpManager* mgr = OpManager::Global(); 98 | std::lock_guard(mgr->mutex); 99 | auto it = mgr->tmap.find(group_name); 100 | if (it != mgr->tmap.end()) { 101 | for (auto& trigger : it->second) { 102 | trigger(this); 103 | } 104 | } 105 | auto& op_group = mgr->op_group; 106 | if (index_ >= op_group.size()) { 107 | op_group.resize(index_ + 1); 108 | } 109 | op_group[index_].insert(group_name); 110 | return *this; 111 | } 112 | 113 | } // namespace nnvm 114 | -------------------------------------------------------------------------------- /src/core/pass.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2016 by Contributors 3 | * \file pass.cc 4 | * \brief Support for pass registry. 5 | */ 6 | #include 7 | #include 8 | 9 | namespace dmlc { 10 | // enable registry 11 | DMLC_REGISTRY_ENABLE(nnvm::PassFunctionReg); 12 | } // namespace dmlc 13 | 14 | namespace nnvm { 15 | 16 | const PassFunctionReg* FindPassDep(const std::string&attr_name) { 17 | for (auto* r : dmlc::Registry::List()) { 18 | for (auto& s : r->graph_attr_targets) { 19 | if (s == attr_name) return r; 20 | } 21 | } 22 | return nullptr; 23 | } 24 | 25 | Graph ApplyPasses(Graph g, 26 | const std::vector& pass) { 27 | std::vector fpass; 28 | for (auto& name : pass) { 29 | auto* reg = dmlc::Registry::Find(name); 30 | CHECK(reg != nullptr) 31 | << "Cannot find pass " << name << " in the registry"; 32 | fpass.push_back(reg); 33 | } 34 | 35 | for (auto r : fpass) { 36 | for (auto& dep : r->graph_attr_dependency) { 37 | if (g.attrs.count(dep) == 0) { 38 | auto* pass_dep = FindPassDep(dep); 39 | std::string msg; 40 | if (pass_dep != nullptr) { 41 | msg = " The attribute is provided by pass " + pass_dep->name; 42 | } 43 | LOG(FATAL) << "Graph attr dependency " << dep 44 | << " is required by pass " << r->name 45 | << " but is not available " 46 | << msg; 47 | } 48 | } 49 | g = r->body(std::move(g)); 50 | } 51 | 52 | return g; 53 | } 54 | 55 | } // namespace nnvm 56 | -------------------------------------------------------------------------------- /src/top/nn/nn_common.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file nn_common.h 4 | * \brief Common utilities for nn ops. 5 | */ 6 | #ifndef NNVM_TOP_NN_NN_COMMON_H_ 7 | #define NNVM_TOP_NN_NN_COMMON_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | namespace nnvm { 19 | namespace top { 20 | 21 | template 22 | inline uint32_t UseBiasNumInputs(const NodeAttrs& attrs) { 23 | const ParamType& param = get(attrs.parsed); 24 | return param.use_bias ? 3 : 2; 25 | } 26 | 27 | template 28 | inline std::vector UseBiasListInputNames(const NodeAttrs& attrs) { 29 | const ParamType& param = nnvm::get(attrs.parsed); 30 | if (param.use_bias) { 31 | return {"data", "weight", "bias"}; 32 | } else { 33 | return {"data", "weight"}; 34 | } 35 | } 36 | 37 | /*! 38 | * \brief Convert shape in src_layout to shape in dst_layout 39 | * \param src original shape 40 | * \param src_layout layout of original shape 41 | * \param dst_layout target layout 42 | * \return shape in target layout 43 | */ 44 | inline TShape ConvertLayout(TShape src, const Layout& src_layout, const Layout& dst_layout) { 45 | if (src_layout == dst_layout) { 46 | return src; 47 | } else if (!src_layout.defined()) { 48 | LOG(FATAL) << "cannot convert undefined layout to " << dst_layout; 49 | } else if (!dst_layout.defined()) { 50 | LOG(FATAL) << "cannot convert " << src_layout << " to undefined layout"; 51 | } 52 | 53 | CHECK(src_layout.convertible(dst_layout)) << "cannot convert from " 54 | << src_layout << " to " << dst_layout; 55 | 56 | TShape dst(dst_layout.ndim()); 57 | for (size_t i = 0; i < src_layout.ndim(); ++i) { 58 | Layout::LayoutDim src_dim = src_layout[i]; 59 | if (Layout::is_superdim(src_dim)) { 60 | int dst_major_pos = dst_layout.indexof(Layout::to_superdim(src_dim)); 61 | int dst_minor_pos = dst_layout.indexof(Layout::to_subdim(src_dim)); 62 | int src_minor_pos = src_layout.indexof(Layout::to_subdim(src_dim)); 63 | int src_factor = src_layout.subsizeof(src_dim); 64 | int dst_factor = dst_layout.subsizeof(src_dim); 65 | 66 | uint32_t src_dim_size = src[i]; 67 | if (src_minor_pos >= 0) { 68 | CHECK_EQ(src_factor, src[src_minor_pos]) << "src shape " << src 69 | << " does not agree with layout " << src_layout; 70 | src_dim_size *= src_factor; 71 | } 72 | 73 | dst[dst_major_pos] = src_dim_size; 74 | if (dst_minor_pos >= 0) { 75 | CHECK_GT(dst_factor, 0); 76 | CHECK_LE(dst_factor, src_dim_size) << "Converting " << src 77 | << " from " << src_layout 78 | << " to " << dst_layout 79 | << ": cannot split dimension size of " 80 | << src_dim_size << " by " << dst_factor; 81 | dst[dst_major_pos] /= dst_factor; 82 | dst[dst_minor_pos] = dst_factor; 83 | } 84 | } 85 | } 86 | return dst; 87 | } 88 | 89 | } // namespace top 90 | } // namespace nnvm 91 | 92 | #endif // NNVM_TOP_NN_NN_COMMON_H_ 93 | -------------------------------------------------------------------------------- /src/top/nn/upsampling.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file pooling.cc 4 | * \brief Property def of pooling operators. 5 | */ 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include "./nn_common.h" 11 | #include "../op_common.h" 12 | #include "../elemwise_op_common.h" 13 | 14 | namespace nnvm { 15 | namespace top { 16 | 17 | DMLC_REGISTER_PARAMETER(UpSamplingParam); 18 | 19 | inline bool UpSamplingInferShape(const nnvm::NodeAttrs& attrs, 20 | std::vector* in_shape, 21 | std::vector* out_shape) { 22 | static const Layout kNCHW("NCHW"); 23 | const UpSamplingParam& param = nnvm::get(attrs.parsed); 24 | CHECK_EQ(in_shape->size(), 1U); 25 | CHECK_EQ(out_shape->size(), 1U); 26 | TShape dshape = (*in_shape)[0]; 27 | if (dshape.ndim() == 0) return false; 28 | dshape = ConvertLayout(dshape, param.layout, kNCHW); 29 | TShape oshape = dshape; 30 | oshape[2] = oshape[2] * param.scale; 31 | oshape[3] = oshape[3] * param.scale; 32 | oshape = ConvertLayout(oshape, kNCHW, param.layout); 33 | NNVM_ASSIGN_OUTPUT_SHAPE(attrs, *out_shape, 0, oshape); 34 | return true; 35 | } 36 | 37 | inline bool UpsamplingLayout(const NodeAttrs& attrs, 38 | std::vector *in_layouts, 39 | const std::vector *last_in_layouts, 40 | std::vector *out_layouts) { 41 | const UpSamplingParam& param = nnvm::get(attrs.parsed); 42 | CHECK_EQ(in_layouts->size(), 1U); 43 | CHECK_EQ(out_layouts->size(), 1U); 44 | const Layout layout(param.layout); 45 | NNVM_ASSIGN_LAYOUT(*in_layouts, 0, layout); 46 | NNVM_ASSIGN_LAYOUT(*out_layouts, 0, layout); 47 | return true; 48 | } 49 | 50 | NNVM_REGISTER_OP(upsampling) 51 | .describe(R"(Perform nearest neighbor upsampling to input array. 52 | 53 | - **data**: Input is 4D array of shape (batch_size, channels, in_height, in_width). 54 | - **out**: Output is 4D array of shape (batch_size, channels, in_height*scale, in_width*scale). 55 | 56 | )" NNVM_ADD_FILELINE) 57 | .add_argument("data", "4D Tensor", "Input data.") 58 | .add_arguments(UpSamplingParam::__FIELDS__()) 59 | .set_attr_parser(ParamParser) 60 | .set_attr("FGetAttrDict", ParamGetAttrDict) 61 | .set_attr("FInferShape", UpSamplingInferShape) 62 | .set_attr("FInferType", ElemwiseType<1, 1>) 63 | .set_attr("FCorrectLayout", UpsamplingLayout) 64 | .set_num_outputs(1) 65 | .set_num_inputs(1) 66 | .set_support_level(2); 67 | 68 | } // namespace top 69 | } // namespace nnvm 70 | -------------------------------------------------------------------------------- /src/top/tensor/state_op.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2018 by Contributors 3 | * \file state_op.cc 4 | * \brief Experimental operators 5 | * Currently we only support assign 6 | */ 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "../op_common.h" 14 | #include "../elemwise_op_common.h" 15 | 16 | namespace nnvm { 17 | namespace top { 18 | 19 | using namespace tvm; 20 | using namespace nnvm::compiler; 21 | 22 | NNVM_REGISTER_OP(_assign) 23 | .describe(R"doc(Assign rhs to the lhs. 24 | 25 | lhs must be a Variable. 26 | This is an experimental operator. 27 | 28 | )doc" NNVM_ADD_FILELINE) 29 | .set_num_inputs(2) 30 | .set_num_outputs(1) 31 | .set_attr( 32 | "FMutateInputs", [](const NodeAttrs& attrs) { 33 | return std::vector{0}; 34 | }) 35 | .set_attr( 36 | "FTVMCompute", [](const NodeAttrs& attrs, 37 | const Array& inputs, 38 | const Array& out_info) { 39 | // This implementation is needed for the special 40 | // logic handling assign in the compiler 41 | // It simply copies the result of rhs the output 42 | // The later decoration in compiler will change 43 | // the memory assignment of assign to tie 44 | // the lhs to the output. 45 | return Array{ topi::identity(inputs[1]) }; 46 | }) 47 | .set_attr("FInferShape", SameShape) 48 | .set_attr( 49 | "FCorrectLayout", [](const NodeAttrs& attrs, 50 | std::vector *in_layouts, 51 | const std::vector *last_in_layouts, 52 | std::vector *out_layouts) { 53 | NNVM_ASSIGN_LAYOUT(*in_layouts, 1, (*in_layouts)[0]); 54 | NNVM_ASSIGN_LAYOUT(*out_layouts, 0, (*in_layouts)[0]); 55 | return true; 56 | }) 57 | .set_attr( 58 | "FInplaceOption", [](const NodeAttrs& attrs) { 59 | return std::vector >{{1, 0}}; 60 | }) 61 | .set_attr( 62 | "FGradient", [](const NodePtr& n, 63 | const std::vector& ograds){ 64 | return std::vector{ 65 | MakeNode("zeros_like", n->attrs.name + "_zero_grad", 66 | {n->inputs[0]}), 67 | ograds[0] 68 | }; 69 | }); 70 | 71 | } // namespace top 72 | } // namespace nnvm 73 | -------------------------------------------------------------------------------- /src/top/vision/yolo2/region.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2018 by Contributors 3 | * \file region.cc 4 | * \brief Property def of pooling operators. 5 | */ 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include "../../op_common.h" 11 | #include "region.h" 12 | 13 | namespace nnvm { 14 | namespace top { 15 | 16 | NNVM_REGISTER_OP(yolo2_region) 17 | .describe(R"code(Region layer 18 | )code" NNVM_ADD_FILELINE) 19 | .set_num_inputs(1) 20 | .set_num_outputs(1) 21 | .set_support_level(5) 22 | .add_argument("data", "Tensor", "Input data") 23 | .set_attr("FInferType", RegionType<1, 1>) 24 | .set_attr("FInferShape", RegionShape<1, 1>) 25 | .set_attr( 26 | "FInplaceOption", 27 | [](const NodeAttrs &attrs) { 28 | return std::vector>{{0, 0}, {1, 0}}; 29 | }) 30 | .set_attr("FGradient", [](const NodePtr &n, 31 | const std::vector &ograds) { 32 | return std::vector{ograds[0], ograds[0]}; 33 | }); 34 | } // namespace top 35 | } // namespace nnvm 36 | -------------------------------------------------------------------------------- /src/top/vision/yolo2/region.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2018 by Contributors 3 | * \file region.h 4 | */ 5 | #ifndef NNVM_TOP_VISION_YOLO2_REGION_H_ 6 | #define NNVM_TOP_VISION_YOLO2_REGION_H_ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace nnvm { 15 | namespace top { 16 | 17 | template 25 | inline bool RegionAttr(const nnvm::NodeAttrs &attrs, 26 | std::vector *in_attrs, 27 | std::vector *out_attrs, 28 | const AttrType &none) { 29 | AttrType dattr = none; 30 | size_t in_size = in_attrs->size(); 31 | size_t out_size = out_attrs->size(); 32 | if (n_in != -1) { 33 | in_size = static_cast(n_in); 34 | } 35 | if (n_out != -1) { 36 | out_size = static_cast(n_out); 37 | } 38 | 39 | auto deduce = [&](std::vector *vec, size_t size, const char *name) { 40 | for (size_t i = 0; i < size; ++i) { 41 | if (i == 0) 42 | CHECK(assign(&dattr, (*vec)[i])) 43 | << "Incompatible attr in node " << attrs.name << " at " << i 44 | << "-th " << name << ": " 45 | << "expected " << attr_string(dattr) << ", got " 46 | << attr_string((*vec)[i]); 47 | } 48 | }; 49 | deduce(in_attrs, in_size, "input"); 50 | 51 | auto write = [&](std::vector *vec, size_t size, const char *name) { 52 | for (size_t i = 0; i < size; ++i) { 53 | CHECK(assign(&(*vec)[i], dattr)) 54 | << "Incompatible attr in node " << attrs.name << " at " << i << "-th " 55 | << name << ": " 56 | << "expected " << attr_string(dattr) << ", got " 57 | << attr_string((*vec)[i]); 58 | } 59 | }; 60 | write(out_attrs, out_size, "output"); 61 | 62 | if (is_none(dattr)) { 63 | return false; 64 | } 65 | return true; 66 | } 67 | 68 | template 69 | inline bool RegionShape(const NodeAttrs &attrs, 70 | std::vector *in_attrs, 71 | std::vector *out_attrs) { 72 | if (n_in != -1) { 73 | CHECK_EQ(in_attrs->size(), static_cast(n_in)) 74 | << " in operator " << attrs.name; 75 | } 76 | if (n_out != -1) { 77 | CHECK_EQ(out_attrs->size(), static_cast(n_out)) 78 | << " in operator " << attrs.name; 79 | } 80 | return RegionAttr( 81 | attrs, in_attrs, out_attrs, TShape()); 82 | } 83 | 84 | template 85 | inline bool RegionType(const NodeAttrs &attrs, 86 | std::vector *in_attrs, 87 | std::vector *out_attrs) { 88 | if (n_in != -1) { 89 | CHECK_EQ(in_attrs->size(), static_cast(n_in)) 90 | << " in operator " << attrs.name; 91 | } 92 | if (n_out != -1) { 93 | CHECK_EQ(out_attrs->size(), static_cast(n_out)) 94 | << " in operator " << attrs.name; 95 | } 96 | return RegionAttr( 97 | attrs, in_attrs, out_attrs, -1); 98 | } 99 | } // namespace top 100 | } // namespace nnvm 101 | #endif // NNVM_TOP_VISION_YOLO2_REGION_H_ 102 | -------------------------------------------------------------------------------- /src/top/vision/yolo2/reorg.cc: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2018 by Contributors 3 | * \file reorg.cc 4 | */ 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include "../../op_common.h" 10 | #include "../../elemwise_op_common.h" 11 | #include "reorg.h" 12 | 13 | namespace nnvm { 14 | namespace top { 15 | 16 | // reorg 17 | DMLC_REGISTER_PARAMETER(ReorgParam); 18 | 19 | inline bool ReorgInferShape(const nnvm::NodeAttrs &attrs, 20 | std::vector *in_shape, 21 | std::vector *out_shape) { 22 | const ReorgParam ¶m = nnvm::get(attrs.parsed); 23 | TShape dshape = in_shape->at(0); 24 | if (dshape.ndim() == 0) 25 | return false; 26 | NNVM_ASSIGN_INPUT_SHAPE(attrs, *in_shape, 0, dshape); 27 | CHECK_EQ(dshape.ndim(), 4) << "Input data should be 4D"; 28 | CHECK_GT(param.stride, 0U) << "Stride value cannot be 0"; 29 | TShape oshape({dshape[0], 0, 0, 0}); 30 | oshape[1] = dshape[1] * param.stride * param.stride; 31 | oshape[2] = dshape[2] / param.stride; 32 | oshape[3] = dshape[3] / param.stride; 33 | NNVM_ASSIGN_OUTPUT_SHAPE(attrs, *out_shape, 0, oshape); 34 | return true; 35 | } 36 | 37 | NNVM_REGISTER_OP(yolo2_reorg) 38 | .describe(R"(Perform reorg operation on input array based on the stride value. 39 | - **data**: Input is 4D array of shape (batch_size, channels, in_height, in_width). 40 | - **out**: Output is 4D array of shape (batch_size, channels/(stride*stride), in_height*stride, in_width*stride). 41 | )" NNVM_ADD_FILELINE) 42 | .set_num_inputs(1) 43 | .set_num_outputs(1) 44 | .set_support_level(5) 45 | .add_argument("data", "Tensor", "Data input to reorganize") 46 | .set_attr_parser(ParamParser) 47 | .add_arguments(ReorgParam::__FIELDS__()) 48 | .set_attr("FGetAttrDict", ParamGetAttrDict) 49 | .set_attr("FInferType", ElemwiseType<-1, 1>) 50 | .set_attr("FInferShape", ReorgInferShape); 51 | } // namespace top 52 | } // namespace nnvm 53 | -------------------------------------------------------------------------------- /src/top/vision/yolo2/reorg.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2018 by Contributors 3 | * \file reorg.h 4 | */ 5 | #ifndef NNVM_TOP_VISION_YOLO2_REORG_H_ 6 | #define NNVM_TOP_VISION_YOLO2_REORG_H_ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | namespace nnvm { 15 | namespace top { 16 | 17 | template 25 | inline bool ReorgAttr(const nnvm::NodeAttrs &attrs, 26 | std::vector *in_attrs, 27 | std::vector *out_attrs, 28 | const AttrType &none) { 29 | AttrType dattr = none; 30 | size_t in_size = in_attrs->size(); 31 | size_t out_size = out_attrs->size(); 32 | if (n_in != -1) { 33 | in_size = static_cast(n_in); 34 | } 35 | if (n_out != -1) { 36 | out_size = static_cast(n_out); 37 | } 38 | 39 | auto deduce = [&](std::vector *vec, size_t size, const char *name) { 40 | for (size_t i = 0; i < size; ++i) { 41 | if (i == 0) { 42 | CHECK(assign(&dattr, (*vec)[i])) 43 | << "Incompatible attr in node " << attrs.name << " at " << i 44 | << "-th " << name << ": " 45 | << "expected " << attr_string(dattr) << ", got " 46 | << attr_string((*vec)[i]); 47 | } 48 | } 49 | }; 50 | deduce(in_attrs, in_size, "input"); 51 | 52 | auto write = [&](std::vector *vec, size_t size, const char *name) { 53 | for (size_t i = 0; i < size; ++i) { 54 | CHECK(assign(&(*vec)[i], dattr)) 55 | << "Incompatible attr in node " << attrs.name << " at " << i << "-th " 56 | << name << ": " 57 | << "expected " << attr_string(dattr) << ", got " 58 | << attr_string((*vec)[i]); 59 | } 60 | }; 61 | write(out_attrs, out_size, "output"); 62 | 63 | if (is_none(dattr)) { 64 | return false; 65 | } 66 | return true; 67 | } 68 | 69 | template 70 | inline bool ReorgShape(const NodeAttrs &attrs, 71 | std::vector *in_attrs, 72 | std::vector *out_attrs) { 73 | if (n_in != -1) { 74 | CHECK_EQ(in_attrs->size(), static_cast(n_in)) 75 | << " in operator " << attrs.name; 76 | } 77 | if (n_out != -1) { 78 | CHECK_EQ(out_attrs->size(), static_cast(n_out)) 79 | << " in operator " << attrs.name; 80 | } 81 | return ReorgAttr( 82 | attrs, in_attrs, out_attrs, TShape()); 83 | } 84 | 85 | template 86 | inline bool ReorgType(const NodeAttrs &attrs, 87 | std::vector *in_attrs, 88 | std::vector *out_attrs) { 89 | if (n_in != -1) { 90 | CHECK_EQ(in_attrs->size(), static_cast(n_in)) 91 | << " in operator " << attrs.name; 92 | } 93 | if (n_out != -1) { 94 | CHECK_EQ(out_attrs->size(), static_cast(n_out)) 95 | << " in operator " << attrs.name; 96 | } 97 | return ReorgAttr( 98 | attrs, in_attrs, out_attrs, -1); 99 | } 100 | 101 | struct ReorgParam : public dmlc::Parameter { 102 | int stride; 103 | 104 | DMLC_DECLARE_PARAMETER(ReorgParam) { 105 | DMLC_DECLARE_FIELD(stride).set_default(1).describe("Stride value"); 106 | } 107 | }; 108 | } // namespace top 109 | } // namespace nnvm 110 | #endif // NNVM_TOP_VISION_YOLO2_REORG_H_ 111 | -------------------------------------------------------------------------------- /tests/ci_build/Dockerfile.gpu: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:8.0-cudnn7-devel 2 | 3 | # Base scripts 4 | RUN apt-get update --fix-missing 5 | 6 | COPY install/ubuntu_install_core.sh /install/ubuntu_install_core.sh 7 | RUN bash /install/ubuntu_install_core.sh 8 | 9 | COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh 10 | RUN bash /install/ubuntu_install_python.sh 11 | 12 | COPY install/ubuntu_install_llvm.sh /install/ubuntu_install_llvm.sh 13 | RUN bash /install/ubuntu_install_llvm.sh 14 | 15 | COPY install/ubuntu_install_opencl.sh /install/ubuntu_install_opencl.sh 16 | RUN bash /install/ubuntu_install_opencl.sh 17 | 18 | COPY install/ubuntu_install_python_package.sh /install/ubuntu_install_python_package.sh 19 | RUN bash /install/ubuntu_install_python_package.sh 20 | 21 | COPY install/ubuntu_install_sphinx.sh /install/ubuntu_install_sphinx.sh 22 | RUN bash /install/ubuntu_install_sphinx.sh 23 | 24 | # Fix recommonmark to latest version 25 | RUN git clone https://github.com/rtfd/recommonmark 26 | RUN cd recommonmark; python setup.py install 27 | 28 | # Enable doxygen for c++ doc build 29 | RUN apt-get update && apt-get install -y doxygen graphviz libprotobuf-dev protobuf-compiler 30 | 31 | # DL Frameworks 32 | COPY install/ubuntu_install_mxnet.sh /install/ubuntu_install_mxnet.sh 33 | RUN bash /install/ubuntu_install_mxnet.sh 34 | 35 | COPY install/ubuntu_install_onnx.sh /install/ubuntu_install_onnx.sh 36 | RUN bash /install/ubuntu_install_onnx.sh 37 | 38 | COPY install/ubuntu_install_coreml.sh /install/ubuntu_install_coreml.sh 39 | RUN bash /install/ubuntu_install_coreml.sh 40 | 41 | COPY install/ubuntu_install_keras.sh /install/ubuntu_install_keras.sh 42 | RUN bash /install/ubuntu_install_keras.sh 43 | 44 | COPY install/ubuntu_install_darknet.sh /install/ubuntu_install_darknet.sh 45 | RUN bash /install/ubuntu_install_darknet.sh 46 | 47 | RUN pip install Pillow 48 | 49 | # Environment variables 50 | ENV PATH=/usr/local/nvidia/bin:${PATH} 51 | ENV PATH=/usr/local/cuda/bin:${PATH} 52 | ENV CPLUS_INCLUDE_PATH=/usr/local/cuda/include:${CPLUS_INCLUDE_PATH} 53 | ENV C_INCLUDE_PATH=/usr/local/cuda/include:${C_INCLUDE_PATH} 54 | ENV LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/nvidia/lib64:${LIBRARY_PATH} 55 | ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH} 56 | -------------------------------------------------------------------------------- /tests/ci_build/Dockerfile.lint: -------------------------------------------------------------------------------- 1 | # For lint test 2 | FROM ubuntu:16.04 3 | 4 | RUN apt-get update && apt-get install -y python-pip sudo 5 | RUN apt-get install -y doxygen graphviz 6 | RUN pip install cpplint pylint 7 | -------------------------------------------------------------------------------- /tests/ci_build/README.md: -------------------------------------------------------------------------------- 1 | # CI Build Scripts 2 | 3 | This directory contains the files and setup instructions to run all tests. 4 | 5 | ## Run locally 6 | 7 | To run locally, we need to first install 8 | [docker](https://docs.docker.com/engine/installation/) and 9 | [nvidia-docker](https://github.com/NVIDIA/nvidia-docker/wiki). 10 | 11 | Then we can run the tasks defined in the [Jenkinsfile](../../Jenkinsfile) by 12 | using (`ci_build.sh`)[./ci_build.sh]. For example 13 | 14 | - lint the python codes 15 | 16 | ```bash 17 | ./ci_build.sh lint make pylint 18 | ``` 19 | 20 | - build codes with CUDA supports 21 | 22 | ```bash 23 | ./ci_build.sh gpu tests/scripts/task_build.sh 24 | ``` 25 | 26 | - do the python unittest 27 | 28 | ```bash 29 | ./ci_build.sh gpu tests/scripts/task_python_test.sh 30 | ``` 31 | 32 | - build the documents. The results will be available at `docs/_build/html` 33 | 34 | ```bash 35 | tests/ci_build/ci_build.sh gpu tests/scripts/task_python_docs.sh 36 | ``` 37 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_core.sh: -------------------------------------------------------------------------------- 1 | # install libraries for building c++ core on ubuntu 2 | apt-get install -y --no-install-recommends --force-yes \ 3 | git make libgtest-dev cmake wget unzip libtinfo-dev libz-dev\ 4 | libcurl4-openssl-dev libopenblas-dev g++ sudo 5 | 6 | cd /usr/src/gtest && cmake CMakeLists.txt && make && cp *.a /usr/lib 7 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_coreml.sh: -------------------------------------------------------------------------------- 1 | pip2 install coremltools 2 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_darknet.sh: -------------------------------------------------------------------------------- 1 | #install the necessary dependancies, cffi, opencv 2 | wget 'https://github.com/siju-samuel/darknet/blob/master/lib/libdarknet.so?raw=true' -O libdarknet.so 3 | pip2 install opencv-python cffi 4 | pip3 install opencv-python cffi 5 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_keras.sh: -------------------------------------------------------------------------------- 1 | pip2 install keras tensorflow h5py 2 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_llvm.sh: -------------------------------------------------------------------------------- 1 | echo deb http://apt.llvm.org/xenial/ llvm-toolchain-xenial-4.0 main\ 2 | >> /etc/apt/sources.list.d/llvm.list 3 | echo deb-src http://apt.llvm.org/xenial/ llvm-toolchain-xenial-4.0 main\ 4 | >> /etc/apt/sources.list.d/llvm.list 5 | 6 | echo deb http://apt.llvm.org/xenial/ llvm-toolchain-xenial-5.0 main\ 7 | >> /etc/apt/sources.list.d/llvm.list 8 | echo deb-src http://apt.llvm.org/xenial/ llvm-toolchain-xenial-5.0 main\ 9 | >> /etc/apt/sources.list.d/llvm.list 10 | 11 | echo deb http://apt.llvm.org/xenial/ llvm-toolchain-xenial main\ 12 | >> /etc/apt/sources.list.d/llvm.list 13 | echo deb-src http://apt.llvm.org/xenial/ llvm-toolchain-xenial main\ 14 | >> /etc/apt/sources.list.d/llvm.list 15 | 16 | wget -O - http://apt.llvm.org/llvm-snapshot.gpg.key|sudo apt-key add - 17 | apt-get update && apt-get install -y --force-yes llvm-4.0 llvm-5.0 llvm-6.0 18 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_mxnet.sh: -------------------------------------------------------------------------------- 1 | pip2 install mxnet 2 | pip3 install mxnet 3 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_onnx.sh: -------------------------------------------------------------------------------- 1 | pip2 install onnx>=1.1.0 2 | pip3 install onnx>=1.1.0 3 | 4 | pip2 install http://download.pytorch.org/whl/cu75/torch-0.2.0.post3-cp27-cp27mu-manylinux1_x86_64.whl 5 | pip2 install torchvision 6 | pip3 install http://download.pytorch.org/whl/cu75/torch-0.2.0.post3-cp35-cp35m-manylinux1_x86_64.whl 7 | pip3 install torchvision 8 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_opencl.sh: -------------------------------------------------------------------------------- 1 | # Install OpenCL runtime in nvidia docker. 2 | apt-get install -y --no-install-recommends --force-yes \ 3 | ocl-icd-libopencl1 \ 4 | clinfo && \ 5 | rm -rf /var/lib/apt/lists/* 6 | 7 | mkdir -p /etc/OpenCL/vendors && \ 8 | echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd 9 | 10 | echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf && \ 11 | echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf 12 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_python.sh: -------------------------------------------------------------------------------- 1 | # install python and pip, don't modify this, modify install_python_package.sh 2 | apt-get update && apt-get install -y python-pip python-dev python3-dev 3 | 4 | # the version of the pip shipped with ubuntu may be too lower, install a recent version here 5 | cd /tmp && wget https://bootstrap.pypa.io/get-pip.py && python3 get-pip.py && python2 get-pip.py 6 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_python_package.sh: -------------------------------------------------------------------------------- 1 | # install libraries for python package on ubuntu 2 | pip2 install nose pylint numpy nose-timer cython decorator scipy tornado 3 | pip3 install nose pylint numpy nose-timer cython decorator scipy tornado 4 | -------------------------------------------------------------------------------- /tests/ci_build/install/ubuntu_install_sphinx.sh: -------------------------------------------------------------------------------- 1 | pip install sphinx==1.6.2 sphinx-gallery sphinx_rtd_theme matplotlib Image commonmark>=0.7.3 docutils>=0.11 2 | -------------------------------------------------------------------------------- /tests/ci_build/with_the_same_user: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # This script is a wrapper creating the same user inside container as the one 4 | # running the ci_build.sh outside the container. It also set the home directory 5 | # for the user inside container to match the same absolute path as the workspace 6 | # outside of container. Do not run this manually. It does not make sense. It is 7 | # intended to be called by ci_build.sh only. 8 | 9 | set -e 10 | 11 | COMMAND=("$@") 12 | 13 | if ! touch /this_is_writable_file_system; then 14 | echo "You can't write to your filesystem!" 15 | echo "If you are in Docker you should check you do not have too many images" \ 16 | "with too many files in them. Docker has some issue with it." 17 | exit 1 18 | else 19 | rm /this_is_writable_file_system 20 | fi 21 | 22 | getent group "${CI_BUILD_GID}" || addgroup --gid "${CI_BUILD_GID}" "${CI_BUILD_GROUP}" 23 | getent passwd "${CI_BUILD_UID}" || adduser --gid "${CI_BUILD_GID}" --uid "${CI_BUILD_UID}" \ 24 | --gecos "${CI_BUILD_USER} (generated by with_the_same_user script)" \ 25 | --disabled-password --home "${CI_BUILD_HOME}" --quiet "${CI_BUILD_USER}" 26 | usermod -a -G sudo "${CI_BUILD_USER}" 27 | echo "${CI_BUILD_USER} ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/90-nopasswd-sudo 28 | 29 | HOME=${CI_BUILD_HOME}\ 30 | sudo -u "#${CI_BUILD_UID}" --preserve-env\ 31 | PATH=${PATH}\ 32 | LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\ 33 | HOME=${CI_BUILD_HOME}\ 34 | ${COMMAND[@]} 35 | -------------------------------------------------------------------------------- /tests/cpp/.gitignore: -------------------------------------------------------------------------------- 1 | unittest 2 | *.d 3 | *_test 4 | -------------------------------------------------------------------------------- /tests/cpp/op_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | NNVM_REGISTER_OP(add) 7 | .describe("add two data together") 8 | .set_num_inputs(2) 9 | .set_attr("inplace_pair", std::make_pair(0, 0)); 10 | 11 | NNVM_REGISTER_OP(add) 12 | .set_attr("nick_name", "plus"); 13 | 14 | 15 | TEST(Op, GetAttr) { 16 | using namespace nnvm; 17 | auto add = Op::Get("add"); 18 | auto nick = Op::GetAttr("nick_name"); 19 | 20 | CHECK_EQ(nick[add], "plus"); 21 | } 22 | 23 | int main(int argc, char ** argv) { 24 | testing::InitGoogleTest(&argc, argv); 25 | testing::FLAGS_gtest_death_test_style = "threadsafe"; 26 | return RUN_ALL_TESTS(); 27 | } 28 | -------------------------------------------------------------------------------- /tests/cpp/tuple_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | TEST(Tuple, Basic) { 6 | using nnvm::Tuple; 7 | using nnvm::TShape; 8 | Tuple x{1, 2, 3}; 9 | Tuple y{1, 2, 3, 5, 6}; 10 | x = std::move(y); 11 | 12 | CHECK_EQ(x.ndim(), 5); 13 | Tuple z{1, 2, 3, 5, 6}; 14 | std::ostringstream os; 15 | os << z; 16 | CHECK_EQ(os.str(), "[1,2,3,5,6]"); 17 | std::istringstream is(os.str()); 18 | is >> y; 19 | CHECK_EQ(x, y); 20 | Tuple ss{1, 2, 3}; 21 | TShape s = ss; 22 | s = std::move(ss); 23 | CHECK((s == TShape{1, 2, 3})); 24 | } 25 | 26 | int main(int argc, char ** argv) { 27 | testing::InitGoogleTest(&argc, argv); 28 | testing::FLAGS_gtest_death_test_style = "threadsafe"; 29 | return RUN_ALL_TESTS(); 30 | } 31 | -------------------------------------------------------------------------------- /tests/cpp/unittest.mk: -------------------------------------------------------------------------------- 1 | GTEST_LIB=$(GTEST_PATH)/lib/ 2 | GTEST_INC=$(GTEST_PATH)/include/ 3 | 4 | TEST_SRC = $(wildcard tests/cpp/*_test.cc) 5 | TEST = $(patsubst tests/cpp/%_test.cc, tests/cpp/%_test, $(TEST_SRC)) 6 | 7 | tests/cpp/%_test: tests/cpp/%_test.cc lib/libnnvm.a 8 | $(CXX) -std=c++11 $(CFLAGS) -MM -MT tests/cpp/$* $< >tests/cpp/$*.d 9 | $(CXX) -std=c++11 $(CFLAGS) -I$(GTEST_INC) -o $@ $(filter %.cc %.a, $^) \ 10 | -L$(GTEST_LIB) $(LDFLAGS) -lgtest 11 | 12 | -include tests/cpp/*.d 13 | -------------------------------------------------------------------------------- /tests/python/compiler/test_alter_op_layout.py: -------------------------------------------------------------------------------- 1 | """Unittest cases for AlterOpLayout pass""" 2 | from nnvm import symbol as sym 3 | from nnvm.compiler import graph_attr 4 | from nnvm.top import registry as reg 5 | import nnvm.graph as graph 6 | 7 | def get_layouts(g): 8 | ldict = {} 9 | vlayout = g.json_attr("layout") 10 | entry_ptr = g.index.entry_ptr 11 | for i, n in enumerate(g.index.nodes): 12 | begin, end = entry_ptr[i], entry_ptr[i + 1] 13 | ldict[n["name"]] = vlayout[begin:end] 14 | return ldict 15 | 16 | 17 | def test_alter_conv2d_layout(): 18 | data = sym.Variable("data", shape=(1, 32, 512, 512)) 19 | conv = sym.conv2d(data, name="conv", channels=16, 20 | kernel_size=(3,3), padding=(1,1), 21 | use_bias=False, layout="NCHW") 22 | # split here 23 | convs = sym.split(conv, indices_or_sections=2) 24 | relus = [sym.relu(x, name="relu") for x in convs] 25 | relu = sym.concatenate(*relus) 26 | flatten = sym.flatten(relu, name="flatten") 27 | softmax = sym.softmax(flatten, name="softmax") 28 | g = graph.create(softmax) 29 | 30 | g = g.apply("CorrectLayout") 31 | g = graph_attr.set_dtype_inputs(g, "float32") 32 | g = g.apply(["InferShape", "InferType"]) 33 | layouts_origin = get_layouts(g) 34 | 35 | @reg.register_alter_op_layout("conv2d", level=100) 36 | def alter_conv2d_layout(attrs, inputs, tinfos): 37 | new_attrs = {k : attrs[k] for k in attrs.keys()} 38 | new_attrs["layout"] = "NCHW16c" 39 | new_attrs["kernel_layout"] = "NCHW16c" 40 | new_attrs["name"] = "conv_alter" 41 | return sym.conv2d(inputs[0], inputs[1], **new_attrs) 42 | 43 | g = g.apply("AlterOpLayout") 44 | layouts = get_layouts(g) 45 | 46 | # check copy layouts 47 | for node in ["data", "relu", "flatten", "softmax", "conv_weight"]: 48 | assert(layouts[node] == layouts_origin[node]) 49 | assert(layouts["conv_alter"] == layouts_origin["conv"]) 50 | 51 | 52 | if __name__ == "__main__": 53 | test_alter_conv2d_layout() 54 | -------------------------------------------------------------------------------- /tests/python/compiler/test_build.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import tvm 4 | from tvm.contrib import graph_runtime 5 | import nnvm.symbol as sym 6 | import nnvm.compiler 7 | from nnvm.compiler.build_module import _run_graph, precompute_prune 8 | 9 | def test_compile(): 10 | x = sym.Variable("x") 11 | y = sym.Variable("y") 12 | z = sym.exp(y + x) 13 | shape = (10, 128) 14 | dtype = tvm.float32 15 | shape_dict = {"x": shape, "y": shape} 16 | def verify(graph, lib): 17 | m = graph_runtime.create(graph, lib, tvm.cpu(0)) 18 | # get member functions 19 | set_input, run, get_output = m["set_input"], m["run"], m["get_output"] 20 | na = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 21 | nb = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 22 | # set inputs 23 | set_input("x", na) 24 | set_input("y", nb) 25 | # execute 26 | run() 27 | # get outputs 28 | out = tvm.nd.empty(shape, dtype) 29 | get_output(0, out) 30 | np.testing.assert_allclose( 31 | out.asnumpy(), np.exp(na.asnumpy() + nb.asnumpy())) 32 | 33 | graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict) 34 | assert graph.index.num_nodes == 3 35 | verify(graph, lib) 36 | 37 | with nnvm.compiler.build_config(opt_level=0): 38 | graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict) 39 | # print(graph.ir()) 40 | assert graph.index.num_nodes == 4 41 | verify(graph, lib) 42 | 43 | def test_run(): 44 | x = sym.Variable("x") 45 | y = sym.Variable("y") 46 | z = sym.exp(y + x) 47 | shape = (10, 10) 48 | dtype = tvm.float32 49 | nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 50 | ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 51 | res = _run_graph(z, {"x": nx, "y": ny}) 52 | np.testing.assert_allclose( 53 | res[0].asnumpy(), np.exp(nx.asnumpy() + ny.asnumpy())) 54 | 55 | 56 | def test_precompute_prune(): 57 | x = sym.Variable("x") + 1 58 | a = sym.Variable("a") 59 | y = sym.Variable("y") 60 | z = y + x + a 61 | shape = (10, 10) 62 | dtype = tvm.float32 63 | nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 64 | na = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 65 | ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 66 | params = {"x": nx, "a": na} 67 | graph, lib, params = nnvm.compiler.build( 68 | z, "llvm", shape={"y": ny.shape}, params=params) 69 | assert graph.index.num_nodes == 4 70 | m = graph_runtime.create(graph, lib, tvm.cpu(0)) 71 | params["y"] = ny 72 | res = tvm.nd.empty(shape) 73 | m["load_params"](nnvm.compiler.save_param_dict(params)) 74 | m.run() 75 | out = m.get_output(0, out=res) 76 | np.testing.assert_allclose( 77 | res.asnumpy(), nx.asnumpy() + 1 + ny.asnumpy() + na.asnumpy()) 78 | 79 | 80 | def test_dtypes(): 81 | x = sym.Variable("x") 82 | y = sym.relu(x) 83 | dshape = (1, 3, 32, 32) 84 | oshape = dshape 85 | for dtype in ['float32', 'float64', 'int32', 'int16', 'int8', 'int64']: 86 | graph, lib, _ = nnvm.compiler.build(y, 'llvm', {"x": dshape}, dtype=dtype) 87 | m = graph_runtime.create(graph, lib, tvm.cpu()) 88 | if 'float' in dtype: 89 | data = np.random.uniform(size=dshape).astype(dtype) 90 | elif 'int' in dtype: 91 | data = np.random.randint(-127, 127, dshape).astype(dtype) 92 | m.run(x=data) 93 | data = (data > 0) * data 94 | out = m.get_output(0, tvm.nd.empty(oshape, dtype)) 95 | np.testing.assert_allclose(out.asnumpy(), data, atol=1e-5, rtol=1e-5) 96 | 97 | 98 | if __name__ == "__main__": 99 | test_precompute_prune() 100 | test_compile() 101 | test_run() 102 | test_dtypes() 103 | -------------------------------------------------------------------------------- /tests/python/compiler/test_compiler_cache.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tvm 3 | from tvm.contrib import graph_runtime 4 | import nnvm.symbol as sym 5 | import nnvm.compiler 6 | 7 | def test_compile_cache(): 8 | x = sym.Variable("x") 9 | y = sym.Variable("y") 10 | z = sym.exp(y + x) 11 | shape = (10, 1) 12 | dtype = tvm.float32 13 | shape_dict = {"x": shape, "y": shape} 14 | def verify(graph, lib): 15 | m = graph_runtime.create(graph, lib, tvm.cpu(0)) 16 | # get member functions 17 | na = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 18 | nb = tvm.nd.array(np.random.uniform(size=shape).astype(dtype)) 19 | m.run(x=na, y=nb) 20 | # get outputs 21 | out = m.get_output(0, tvm.nd.empty(shape, dtype)) 22 | np.testing.assert_allclose( 23 | out.asnumpy(), np.exp(na.asnumpy() + nb.asnumpy())) 24 | 25 | engine = nnvm.compiler.engine 26 | graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict) 27 | inputs = [tvm.placeholder((10,)), tvm.placeholder((10,))] 28 | 29 | gkey = nnvm.compiler.graph_key(nnvm.graph.create(z), inputs, "llvm") 30 | gkey2 = nnvm.compiler.graph_key(nnvm.graph.create(z), inputs + inputs, "llvm") 31 | gf = engine[gkey] 32 | assert gf is not None 33 | assert engine[gkey2] is None 34 | graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict) 35 | assert graph.index.num_nodes == 3 36 | verify(graph, lib) 37 | # Test various set external cache 38 | engine.clear_cache() 39 | engine[gkey] = gf 40 | 41 | if __name__ == "__main__": 42 | test_compile_cache() 43 | -------------------------------------------------------------------------------- /tests/python/compiler/test_graph_pass.py: -------------------------------------------------------------------------------- 1 | """Unittest cases for graph pass""" 2 | import nnvm 3 | import nnvm.compiler 4 | from nnvm import symbol as sym 5 | from nnvm.compiler import graph_util, graph_attr 6 | 7 | def test_infer_attr(): 8 | x = sym.Variable("x") 9 | y = x * 2 10 | g = nnvm.graph.create(y) 11 | ishape, oshape = graph_util.infer_shape(g, x=(10,20)) 12 | assert tuple(oshape[0]) == (10, 20) 13 | 14 | itype, otype = graph_util.infer_dtype(g, x="float32") 15 | assert otype[0] == "float32" 16 | 17 | if __name__ == "__main__": 18 | test_infer_attr() 19 | -------------------------------------------------------------------------------- /tests/python/compiler/test_nhwc_layout.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tvm 3 | from tvm.contrib import graph_runtime as runtime 4 | import nnvm.symbol as sym 5 | import nnvm.compiler 6 | from nnvm.testing.config import ctx_list 7 | 8 | def get_sym(layout, kernel_layout, channels): 9 | data = sym.Variable(name="data") 10 | data = sym.conv2d(data=data, kernel_size=(3,3), channels=channels, padding=(1, 1), 11 | layout=layout, kernel_layout=kernel_layout, use_bias=True) 12 | data = sym.max_pool2d(data=data, pool_size=(2, 2), strides=(2, 2), layout=layout) 13 | data = sym.upsampling(data=data, scale=2, layout=layout) 14 | softmax_axis = 1 15 | if layout == "NHWC": 16 | softmax_axis = 3 17 | data = sym.softmax(data=data, axis=softmax_axis) 18 | return data 19 | 20 | 21 | def build_and_run(sym, params, data, out_shape): 22 | ctx = tvm.cpu(0) 23 | graph, lib, params = nnvm.compiler.build(sym, "llvm", shape={"data":data.shape}, params=params) 24 | module = runtime.create(graph, lib, ctx) 25 | module.set_input(**params) 26 | module.set_input("data", data) 27 | module.run() 28 | out = module.get_output(0, tvm.nd.empty(out_shape)) 29 | return out.asnumpy() 30 | 31 | 32 | def test_nhwc(): 33 | data_shape = (1, 3, 224, 224) 34 | out_channel = 8 35 | nchw_sym = get_sym("NCHW", "OIHW", out_channel) 36 | nhwc_sym = get_sym("NHWC", "HWIO", out_channel) 37 | conv_weight = np.random.uniform(-1, 1, (out_channel, 3, 3, 3)).astype(np.float32) 38 | conv_bias = np.random.uniform(-1, 1, (out_channel)).astype(np.float32) 39 | nchw_params = { 40 | "conv2d0_weight" : tvm.nd.array(conv_weight, ctx=tvm.cpu(0)), 41 | "conv2d0_bias" : tvm.nd.array(conv_bias, ctx=tvm.cpu(0)) 42 | } 43 | nhwc_params = { 44 | "conv2d1_weight" : tvm.nd.array(conv_weight.transpose(2, 3, 1, 0), ctx=tvm.cpu(0)), 45 | "conv2d1_bias" : tvm.nd.array(conv_bias, ctx=tvm.cpu(0)) 46 | } 47 | 48 | data = np.random.uniform(-1, 1, data_shape).astype(np.float32) 49 | oshape = (1, out_channel, 224, 224) 50 | oshape_nhwc = (1, 224, 224, out_channel) 51 | nchw_output = build_and_run(nchw_sym, nchw_params, data, oshape) 52 | nhwc_output = build_and_run(nhwc_sym, nhwc_params, data.transpose(0, 2, 3, 1), oshape_nhwc) 53 | np.testing.assert_allclose(nchw_output, nhwc_output.transpose(0, 3, 1, 2), rtol=1e-5, atol=1e-5) 54 | 55 | 56 | if __name__ == "__main__": 57 | test_nhwc() 58 | -------------------------------------------------------------------------------- /tests/python/compiler/test_op_fusion.py: -------------------------------------------------------------------------------- 1 | import nnvm 2 | import numpy as np 3 | import tvm 4 | import topi.testing 5 | from tvm.contrib import graph_runtime 6 | from nnvm import symbol as sym 7 | from nnvm.compiler import graph_util, graph_attr 8 | from nnvm.testing import ctx_list 9 | 10 | def test_ewise_injective(): 11 | x = sym.Variable("x") 12 | y = x * 2 13 | y = sym.flatten(y) + 1 14 | dshape = (10, 2, 3) 15 | shape_dict = {"x": dshape} 16 | dtype = "float32" 17 | target = "llvm" 18 | for target, ctx in ctx_list(): 19 | graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) 20 | assert graph.index.num_nodes == 2 21 | m = graph_runtime.create(graph, lib, ctx) 22 | x_np = np.random.uniform(size=dshape).astype(dtype) 23 | m.run(x=x_np) 24 | out = m.get_output(0, tvm.nd.empty((10, 6))) 25 | np.testing.assert_allclose( 26 | out.asnumpy(), x_np.reshape(out.shape) * 2 + 1, 27 | atol=1e-5, rtol=1e-5) 28 | 29 | 30 | def test_conv_ewise_injective(): 31 | x = sym.Variable("x") 32 | y = sym.conv2d(x, channels=32, kernel_size=(3, 3), groups=32, 33 | name="y", padding=(1,1)) 34 | y = sym.flatten(y + 1) + 1 35 | dtype = "float32" 36 | dshape = (1, 32, 18, 18) 37 | kshape = (32, 1, 3, 3) 38 | oshape = (1, 32* 18 * 18) 39 | shape_dict = {"x": dshape} 40 | 41 | for target, ctx in ctx_list(): 42 | graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) 43 | m = graph_runtime.create(graph, lib, ctx) 44 | # print(graph.ir(join_entry_attrs=["shape"])) 45 | assert graph.index.num_nodes == 5 46 | # set input 47 | data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) 48 | kernel = tvm.nd.array(np.random.uniform(size=kshape).astype(dtype)) 49 | bias = tvm.nd.array(np.random.uniform(size=kshape[0]).astype(dtype)) 50 | m.run(x=data, y_weight=kernel, y_bias=bias) 51 | # get output 52 | out = m.get_output(0, tvm.nd.empty(oshape, dtype)) 53 | c_np = topi.testing.depthwise_conv2d_python_nchw( 54 | data.asnumpy(), kernel.asnumpy(), (1,1), 'SAME') 55 | c_np = c_np + bias.asnumpy().reshape(kshape[0], 1, 1) + 1 56 | c_np = c_np.reshape(c_np.shape[0], np.prod(c_np.shape[1:])) + 1 57 | np.testing.assert_allclose(out.asnumpy(), c_np, rtol=1e-5) 58 | 59 | 60 | def test_injective_reduce_injective(): 61 | x = sym.Variable("x") 62 | x = sym.flatten(x) + 1 63 | y = sym.sum(x, axis=1) 64 | dtype = "float32" 65 | dshape = (32, 1, 18, 18) 66 | shape_dict = {"x": dshape} 67 | 68 | for target, ctx in ctx_list(): 69 | graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) 70 | m = graph_runtime.create(graph, lib, ctx) 71 | assert graph.index.num_nodes == 2 72 | data = np.random.uniform(size=dshape).astype(dtype) 73 | m.run(x=data) 74 | c_np = np.sum(data.reshape(32, 18 * 18) + 1, axis=1) 75 | # get output 76 | out = m.get_output(0, tvm.nd.empty(c_np.shape, dtype)) 77 | np.testing.assert_allclose(out.asnumpy(), c_np, rtol=1e-5) 78 | 79 | 80 | if __name__ == "__main__": 81 | test_injective_reduce_injective() 82 | test_ewise_injective() 83 | test_conv_ewise_injective() 84 | -------------------------------------------------------------------------------- /tests/python/compiler/test_param_dict.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import nnvm.compiler 3 | 4 | def test_save_load(): 5 | x = np.random.uniform(size=(10, 2)).astype("float32") 6 | y = np.random.uniform(size=(1, 2, 3)).astype("float32") 7 | x[:] = 1 8 | y[:] = 1 9 | params = {"x": x, "y": y} 10 | param_bytes = nnvm.compiler.save_param_dict(params) 11 | assert isinstance(param_bytes, bytearray) 12 | param2 = nnvm.compiler.load_param_dict(param_bytes) 13 | assert len(param2) == 2 14 | np.testing.assert_equal(param2["x"].asnumpy(), x) 15 | np.testing.assert_equal(param2["y"].asnumpy(), y) 16 | 17 | 18 | if __name__ == "__main__": 19 | test_save_load() 20 | -------------------------------------------------------------------------------- /tests/python/compiler/test_rpc_exec.py: -------------------------------------------------------------------------------- 1 | import tvm 2 | from tvm.contrib import util, rpc, graph_runtime 3 | import nnvm.symbol as sym 4 | import nnvm.compiler 5 | import numpy as np 6 | 7 | def test_rpc_executor(): 8 | host = "localhost" 9 | port = 9100 10 | server = rpc.Server(host, port, use_popen=True) 11 | 12 | x = sym.Variable("x") 13 | y = sym.Variable("y") 14 | z = sym.exp(y + x) 15 | shape = (10, 128) 16 | dtype = tvm.float32 17 | shape_dict = {"x": shape, "y": shape} 18 | tmp = util.tempdir() 19 | lib_name = tmp.relpath("net.o") 20 | 21 | graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict) 22 | # save module 23 | lib.save(lib_name) 24 | remote = rpc.connect(host, port) 25 | remote.upload(lib_name) 26 | ctx = remote.cpu(0) 27 | # load remote 28 | rlib = remote.load_module("net.o") 29 | 30 | # Create remotemodule 31 | m = graph_runtime.create(graph, rlib, remote.cpu(0)) 32 | # get member functions 33 | set_input, run, get_output = m["set_input"], m["run"], m["get_output"] 34 | na = tvm.nd.array(np.ones(shape).astype(dtype), ctx) 35 | nb = tvm.nd.array(np.ones(shape).astype(dtype), ctx) 36 | # set inputs 37 | set_input("x", na) 38 | set_input("y", nb) 39 | # execute 40 | run() 41 | # get outputs 42 | out = tvm.nd.empty(shape, dtype, ctx) 43 | get_output(0, out) 44 | np.testing.assert_allclose( 45 | out.asnumpy(), np.exp(na.asnumpy() + nb.asnumpy())) 46 | server.terminate() 47 | 48 | if __name__ == "__main__": 49 | test_rpc_executor() 50 | -------------------------------------------------------------------------------- /tests/python/compiler/test_simplify_inference.py: -------------------------------------------------------------------------------- 1 | """Unittest cases for simplify batch_norm""" 2 | import nnvm 3 | from nnvm import symbol as sym 4 | from nnvm.compiler import graph_util, graph_attr 5 | 6 | def test_simplify_batchnorm(): 7 | def simple_bn(x, gamma, beta, moving_mean, moving_var, 8 | axis=1, epsilon=1e-5, shape=None): 9 | # expect = (x - moving_mean) / sym.sqrt(moving_var + eps) * gamma + beta 10 | scale = sym.elemwise_mul(1 / sym.sqrt(moving_var + epsilon), gamma) 11 | shift = sym.elemwise_add( 12 | sym.elemwise_mul(sym.negative(moving_mean), scale), beta) 13 | shape = [-1 if i == axis else 1 for i in range(len(shape))] 14 | # for 2D 15 | num_newaxis=len(shape) - axis - 1 16 | if num_newaxis: 17 | scale = sym.expand_dims(scale, axis=1, num_newaxis=num_newaxis) 18 | shift = sym.expand_dims(shift, axis=1, num_newaxis=num_newaxis) 19 | return x * scale + shift 20 | 21 | 22 | # Before simplify 23 | def check(dim, axis, nstep): 24 | eps = 0.01 25 | x = sym.Variable("x") + 1 26 | beta = sym.Variable("beta") 27 | gamma = sym.Variable("gamma") 28 | moving_var = sym.Variable("moving_var") 29 | moving_mean = sym.Variable("moving_mean") 30 | y1, y2 = x, sym.Variable("xx") + 1 31 | ishape = {"x": tuple(10 for i in range(dim))} 32 | for i in range(nstep): 33 | y1 = sym.batch_norm( 34 | y1 + 1, gamma, beta, moving_mean, moving_var, epsilon=eps, axis=axis) 35 | y1 = sym.dropout(y1) 36 | y2 = simple_bn(y2 + 1, gamma, beta, moving_mean, moving_var, 37 | epsilon=eps, axis=axis, shape=ishape["x"]) 38 | g = nnvm.graph.create(y1) 39 | g2 = nnvm.graph.create(y2) 40 | graph_attr.set_shape_inputs(g, ishape) 41 | g1 = g.apply("InferShape").apply("SimplifyInference") 42 | # assert graph equals as expected 43 | graph_util.check_graph_equal(g1, g2) 44 | 45 | check(2, 1, 1) 46 | check(4, 0, 3) 47 | check(4, 1, 2) 48 | 49 | if __name__ == "__main__": 50 | test_simplify_batchnorm() 51 | -------------------------------------------------------------------------------- /tests/python/compiler/test_top_assign.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import tvm 4 | from tvm.contrib import graph_runtime 5 | 6 | import nnvm.symbol as sym 7 | import nnvm.compiler 8 | from nnvm.testing.config import ctx_list 9 | 10 | 11 | def test_update(): 12 | w = sym.Variable("w") 13 | w2 = sym.Variable("w2") 14 | w = sym._assign(w, w + 1) 15 | w2 = sym._assign(w2, w + 1) 16 | 17 | dshape = (5, 3, 18, 18) 18 | shape_dict = {"w": dshape, "w2":dshape} 19 | dtype = "float32" 20 | 21 | def check(target, ctx): 22 | graph, lib, _ = nnvm.compiler.build(w2, target, shape_dict) 23 | 24 | m = graph_runtime.create(graph, lib, ctx) 25 | 26 | data = tvm.nd.array(np.random.uniform(size=dshape).astype(dtype)) 27 | m.set_input("w", data) 28 | m.run() 29 | out = m.get_input("w2", tvm.nd.empty(dshape, dtype)) 30 | np.testing.assert_allclose(out.asnumpy(), data.asnumpy() + 2, rtol=1e-5) 31 | 32 | m.run() 33 | out = m.get_input("w2", tvm.nd.empty(dshape, dtype)) 34 | np.testing.assert_allclose(out.asnumpy(), data.asnumpy() + 3, rtol=1e-5) 35 | 36 | for target, ctx in ctx_list(): 37 | check(target, ctx) 38 | 39 | 40 | if __name__ == "__main__": 41 | test_update() 42 | -------------------------------------------------------------------------------- /tests/python/frontend/coreml/model_zoo/.gitignore: -------------------------------------------------------------------------------- 1 | *.mlmodel 2 | *.jpg 3 | *.png 4 | -------------------------------------------------------------------------------- /tests/python/frontend/coreml/model_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | import urllib 2 | import os 3 | from PIL import Image 4 | import numpy as np 5 | 6 | def download(url, path, overwrite=False): 7 | if os.path.exists(path) and not overwrite: 8 | return 9 | print('Downloading {} to {}.'.format(url, path)) 10 | urllib.URLopener().retrieve(url, path) 11 | 12 | def get_mobilenet(): 13 | url = 'https://docs-assets.developer.apple.com/coreml/models/MobileNet.mlmodel' 14 | dst = 'mobilenet.mlmodel' 15 | real_dst = os.path.abspath(os.path.join(os.path.dirname(__file__), dst)) 16 | download(url, real_dst) 17 | return os.path.abspath(real_dst) 18 | 19 | def get_resnet50(): 20 | url = 'https://docs-assets.developer.apple.com/coreml/models/Resnet50.mlmodel' 21 | dst = 'resnet50.mlmodel' 22 | real_dst = os.path.abspath(os.path.join(os.path.dirname(__file__), dst)) 23 | download(url, real_dst) 24 | return os.path.abspath(real_dst) 25 | 26 | def get_cat_image(): 27 | url = 'https://gist.githubusercontent.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/fa7ef0e9c9a5daea686d6473a62aacd1a5885849/cat.png' 28 | dst = 'cat.jpg' 29 | real_dst = os.path.abspath(os.path.join(os.path.dirname(__file__), dst)) 30 | download(url, real_dst) 31 | img = Image.open(real_dst).resize((224, 224)) 32 | img = np.transpose(img, (2, 0, 1))[np.newaxis, :] 33 | return np.asarray(img) 34 | -------------------------------------------------------------------------------- /tests/python/frontend/coreml/test_forward.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import topi 4 | import tvm 5 | from tvm.contrib import graph_runtime 6 | import nnvm.symbol as sym 7 | import nnvm.compiler 8 | from nnvm.testing.config import ctx_list 9 | from nnvm import frontend 10 | import coremltools as cm 11 | import model_zoo 12 | 13 | def get_tvm_output(symbol, x, params, target, ctx, 14 | out_shape=(1000,), input_name='image', dtype='float32'): 15 | shape_dict = {input_name : x.shape} 16 | with nnvm.compiler.build_config(opt_level=3): 17 | graph, lib, params = nnvm.compiler.build(symbol, target, shape_dict, params=params) 18 | m = graph_runtime.create(graph, lib, ctx) 19 | # set inputs 20 | m.set_input(input_name, tvm.nd.array(x.astype(dtype))) 21 | m.set_input(**params) 22 | m.run() 23 | # get outputs 24 | out = m.get_output(0, tvm.nd.empty(out_shape, dtype)) 25 | return out.asnumpy() 26 | 27 | def test_model_checkonly(model_file, model_name=''): 28 | model = cm.models.MLModel(model_file) 29 | sym, params = nnvm.frontend.from_coreml(model) 30 | x = model_zoo.get_cat_image() 31 | for target, ctx in ctx_list(): 32 | tvm_output = get_tvm_output(sym, x, params, target, ctx) 33 | print(target, ctx, model_name, 'prediction id: ', np.argmax(tvm_output.flat)) 34 | 35 | def test_mobilenet_checkonly(): 36 | model_file = model_zoo.get_mobilenet() 37 | test_model_checkonly(model_file, 'mobilenet') 38 | 39 | def test_resnet50_checkonly(): 40 | model_file = model_zoo.get_resnet50() 41 | test_model_checkonly(model_file, 'resnet50') 42 | 43 | if __name__ == '__main__': 44 | test_mobilenet_checkonly() 45 | test_resnet50_checkonly() 46 | -------------------------------------------------------------------------------- /tests/python/frontend/mxnet/model_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | """MXNet and NNVM model zoo.""" 2 | from __future__ import absolute_import 3 | from . import mlp, resnet, vgg 4 | import nnvm.testing 5 | 6 | __all__ = ['mx_mlp', 'nnvm_mlp', 'mx_resnet', 'nnvm_resnet', 'mx_vgg', 'nnvm_vgg'] 7 | 8 | _num_class = 1000 9 | 10 | # mlp fc 11 | mx_mlp = mlp.get_symbol(_num_class) 12 | nnvm_mlp = nnvm.testing.mlp.get_workload(1, _num_class)[0] 13 | 14 | # resnet fc 15 | mx_resnet = {} 16 | nnvm_resnet = {} 17 | for num_layer in [18, 34, 50, 101, 152, 200, 269]: 18 | mx_resnet[num_layer] = resnet.get_symbol(_num_class, num_layer, '3,224,224') 19 | nnvm_resnet[num_layer] = nnvm.testing.resnet.get_workload( 20 | 1, _num_class, num_layers=num_layer)[0] 21 | 22 | # vgg fc 23 | mx_vgg = {} 24 | nnvm_vgg = {} 25 | for num_layer in [11, 13, 16, 19]: 26 | mx_vgg[num_layer] = vgg.get_symbol(_num_class, num_layer) 27 | nnvm_vgg[num_layer] = nnvm.testing.vgg.get_workload( 28 | 1, _num_class, num_layers=num_layer)[0] 29 | -------------------------------------------------------------------------------- /tests/python/frontend/mxnet/model_zoo/mlp.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | 18 | """ 19 | a simple multilayer perceptron 20 | """ 21 | import mxnet as mx 22 | 23 | def get_symbol(num_classes=10, **kwargs): 24 | data = mx.symbol.Variable('data') 25 | data = mx.sym.Flatten(data=data) 26 | try: 27 | fc1 = mx.symbol.FullyConnected(data = data, name='fc1', num_hidden=128, flatten=False) 28 | act1 = mx.symbol.Activation(data = fc1, name='relu1', act_type="relu") 29 | fc2 = mx.symbol.FullyConnected(data = act1, name = 'fc2', num_hidden = 64, flatten=False) 30 | act2 = mx.symbol.Activation(data = fc2, name='relu2', act_type="relu") 31 | fc3 = mx.symbol.FullyConnected(data = act2, name='fc3', num_hidden=num_classes, flatten=False) 32 | mlp = mx.symbol.softmax(data = fc3, name = 'softmax') 33 | except: 34 | fc1 = mx.symbol.FullyConnected(data = data, name='fc1', num_hidden=128) 35 | act1 = mx.symbol.Activation(data = fc1, name='relu1', act_type="relu") 36 | fc2 = mx.symbol.FullyConnected(data = act1, name = 'fc2', num_hidden = 64) 37 | act2 = mx.symbol.Activation(data = fc2, name='relu2', act_type="relu") 38 | fc3 = mx.symbol.FullyConnected(data = act2, name='fc3', num_hidden=num_classes) 39 | mlp = mx.symbol.softmax(data = fc3, name = 'softmax') 40 | return mlp 41 | -------------------------------------------------------------------------------- /tests/python/frontend/mxnet/test_graph.py: -------------------------------------------------------------------------------- 1 | import mxnet as mx 2 | import nnvm 3 | from nnvm.compiler import graph_util, graph_attr 4 | import model_zoo 5 | 6 | def compare_graph(sym1, sym2, ishape=(2, 3, 224, 224)): 7 | g1 = nnvm.graph.create(sym1) 8 | g2 = nnvm.graph.create(sym2) 9 | graph_attr.set_shape_inputs(g1, {'data':ishape}) 10 | graph_attr.set_shape_inputs(g2, {'data':ishape}) 11 | g1 = g1.apply("InferShape").apply("SimplifyInference") 12 | g2 = g2.apply("InferShape").apply("SimplifyInference") 13 | graph_util.check_graph_equal(g1, g2) 14 | 15 | def test_mlp(): 16 | mx_sym = model_zoo.mx_mlp 17 | from_mx_sym, _ = nnvm.frontend.from_mxnet(mx_sym) 18 | nnvm_sym = model_zoo.nnvm_mlp 19 | compare_graph(from_mx_sym, nnvm_sym) 20 | 21 | def test_vgg(): 22 | for n in [11, 13, 16, 19]: 23 | mx_sym = model_zoo.mx_vgg[n] 24 | from_mx_sym, _ = nnvm.frontend.from_mxnet(mx_sym) 25 | nnvm_sym = model_zoo.nnvm_vgg[n] 26 | compare_graph(from_mx_sym, nnvm_sym) 27 | 28 | def test_resnet(): 29 | for n in [18, 34, 50, 101]: 30 | mx_sym = model_zoo.mx_resnet[n] 31 | from_mx_sym, _ = nnvm.frontend.from_mxnet(mx_sym) 32 | nnvm_sym = model_zoo.nnvm_resnet[n] 33 | compare_graph(from_mx_sym, nnvm_sym) 34 | 35 | def test_multi_outputs(): 36 | def compose(F, **kwargs): 37 | x = F.sym.Variable('x') 38 | y = F.sym.Variable('y') 39 | z = F.sym.split(x, **kwargs) 40 | return F.sym.broadcast_sub(F.sym.broadcast_add(z[0], z[2]), y) 41 | mx_sym = compose(mx, num_outputs=3, axis=1) 42 | from_mx_sym, _ = nnvm.frontend.from_mxnet(mx_sym) 43 | nnvm_sym = compose(nnvm, indices_or_sections=3, axis=1) 44 | compare_graph(from_mx_sym, nnvm_sym) 45 | 46 | if __name__ == '__main__': 47 | test_mlp() 48 | test_vgg() 49 | test_resnet() 50 | test_multi_outputs() 51 | -------------------------------------------------------------------------------- /tests/python/frontend/onnx/model_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | """Store for onnx examples and common models.""" 2 | from __future__ import absolute_import as _abs 3 | import os 4 | import logging 5 | from .super_resolution import get_super_resolution 6 | 7 | def _download(url, filename, overwrite=False): 8 | if os.path.isfile(filename) and not overwrite: 9 | logging.debug('File %s existed, skip.', filename) 10 | return 11 | logging.debug('Downloading from url %s to %s', url, filename) 12 | try: 13 | import urllib.request 14 | urllib.request.urlretrieve(url, filename) 15 | except: 16 | import urllib 17 | urllib.urlretrieve(url, filename) 18 | 19 | def _as_abs_path(fname): 20 | cur_dir = os.path.abspath(os.path.dirname(__file__)) 21 | return os.path.join(cur_dir, fname) 22 | 23 | 24 | URLS = { 25 | 'super_resolution.onnx': 'https://gist.github.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/93672b029103648953c4e5ad3ac3aadf346a4cdc/super_resolution_0.2.onnx', 26 | 'squeezenet1_1.onnx': 'https://gist.github.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/93672b029103648953c4e5ad3ac3aadf346a4cdc/squeezenet1_1_0.2.onnx', 27 | 'lenet.onnx': 'https://gist.github.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/93672b029103648953c4e5ad3ac3aadf346a4cdc/lenet_0.2.onnx', 28 | 'resnet18_1_0.onnx': 'https://gist.github.com/zhreshold/bcda4716699ac97ea44f791c24310193/raw/b385b1b242dc89a35dd808235b885ed8a19aedc1/resnet18_1.0.onnx'} 29 | 30 | # download and add paths 31 | for k, v in URLS.items(): 32 | name = k.split('.')[0] 33 | path = _as_abs_path(k) 34 | _download(v, path, False) 35 | locals()[name] = path 36 | 37 | # symbol for graph comparison 38 | super_resolution_sym = get_super_resolution() 39 | -------------------------------------------------------------------------------- /tests/python/frontend/onnx/model_zoo/super_resolution.py: -------------------------------------------------------------------------------- 1 | """NNVM symbol corresponding to super_resolution.onnx example.""" 2 | from nnvm import sym 3 | 4 | def get_super_resolution(): 5 | factor = 3 6 | size = 224 7 | data = sym.Variable(name='9') 8 | conv1 = sym.conv2d(data, channels=64, kernel_size=(5, 5), padding=(2, 2), use_bias=False) 9 | relu1 = sym.relu(conv1 + sym.expand_dims(sym.Variable(name='2', shape=(64)), axis=1, num_newaxis=2)) 10 | conv2 = sym.conv2d(relu1, channels=64, kernel_size=(3, 3), padding=(1, 1), use_bias=False) 11 | relu2 = sym.relu(conv2 + sym.expand_dims(sym.Variable(name='4', shape=(64)), axis=1, num_newaxis=2)) 12 | conv3 = sym.conv2d(relu2, channels=32, kernel_size=(3, 3), padding=(1, 1), use_bias=False) 13 | relu3 = sym.relu(conv3 + sym.expand_dims(sym.Variable(name='6', shape=(32)), axis=1, num_newaxis=2)) 14 | conv4 = sym.conv2d(relu3, channels=factor**2, kernel_size=(3, 3), padding=(1, 1), use_bias=False) 15 | conv4 = conv4 + sym.expand_dims(sym.Variable(name='8', shape=(factor**2)), axis=1, num_newaxis=2) 16 | # TODO(zhreshold): allow shape inference for batch size > 1 17 | r1 = sym.reshape(conv4, shape=(1, 1, factor, factor, size, size)) 18 | t1 = sym.transpose(r1, axes=(0, 1, 4, 2, 5, 3)) 19 | r2 = sym.reshape(t1, shape=(1, 1, size * factor, size * factor)) 20 | return r2 21 | -------------------------------------------------------------------------------- /tests/python/frontend/onnx/test_forward.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import nnvm 3 | import tvm 4 | from tvm.contrib import graph_runtime 5 | from nnvm.testing.config import ctx_list 6 | import onnx 7 | from model_zoo import super_resolution, squeezenet1_1, lenet, resnet18_1_0 8 | 9 | def verify_onnx_forward_impl(graph_file, data_shape, out_shape): 10 | import caffe2.python.onnx.backend 11 | def get_caffe2_output(model, x, dtype='float32'): 12 | prepared_backend = caffe2.python.onnx.backend.prepare(model) 13 | W = {model.graph.input[0].name: x.astype(dtype)} 14 | c2_out = prepared_backend.run(W)[0] 15 | return c2_out 16 | 17 | def get_tvm_output(model, x, target, ctx, dtype='float32'): 18 | new_sym, params = nnvm.frontend.from_onnx(model) 19 | shape_dict = {'input_0': x.shape} 20 | graph, lib, params = nnvm.compiler.build(new_sym, target, shape_dict, params=params) 21 | m = graph_runtime.create(graph, lib, ctx) 22 | # set inputs 23 | m.set_input('input_0', tvm.nd.array(x.astype(dtype))) 24 | m.set_input(**params) 25 | m.run() 26 | # get outputs 27 | out = m.get_output(0, tvm.nd.empty(out_shape, dtype)) 28 | return out.asnumpy() 29 | 30 | dtype = 'float32' 31 | x = np.random.uniform(size=data_shape) 32 | model = onnx.load(graph_file) 33 | c2_out = get_caffe2_output(model, x, dtype) 34 | for target, ctx in ctx_list(): 35 | tvm_out = get_tvm_output(model, x, target, ctx, dtype) 36 | np.testing.assert_allclose(c2_out, tvm_out, rtol=1e-5, atol=1e-5) 37 | 38 | def verify_super_resolution_example(): 39 | verify_onnx_forward_impl(super_resolution, (1, 1, 224, 224), (1, 1, 672, 672)) 40 | 41 | def verify_squeezenet1_1(): 42 | verify_onnx_forward_impl(squeezenet1_1, (1, 3, 224, 224), (1, 1000)) 43 | 44 | def verify_lenet(): 45 | verify_onnx_forward_impl(lenet, (1, 1, 28, 28), (1, 10)) 46 | 47 | def verify_resnet18(): 48 | verify_onnx_forward_impl(resnet18_1_0, (1, 3, 224, 224), (1, 1000)) 49 | 50 | if __name__ == '__main__': 51 | # verify_super_resolution_example() 52 | # verify_squeezenet1_1() 53 | # verify_lenet() 54 | verify_resnet18() 55 | -------------------------------------------------------------------------------- /tests/python/frontend/onnx/test_graph.py: -------------------------------------------------------------------------------- 1 | """Test graph equality of onnx models.""" 2 | import nnvm 3 | import onnx 4 | from nnvm.compiler import graph_util, graph_attr 5 | from model_zoo import super_resolution, super_resolution_sym 6 | 7 | def compare_graph(onnx_file, nnvm_sym, ishape): 8 | onnx_model = onnx.load(onnx_file) 9 | onnx_sym, params = nnvm.frontend.from_onnx(onnx_model) 10 | g1 = nnvm.graph.create(onnx_sym) 11 | g2 = nnvm.graph.create(nnvm_sym) 12 | ishapes = {'input_0': ishape} 13 | graph_attr.set_shape_inputs(g1, ishapes) 14 | graph_attr.set_shape_inputs(g2, ishapes) 15 | g1 = g1.apply("InferShape").apply("SimplifyInference") 16 | g2 = g2.apply("InferShape").apply("SimplifyInference") 17 | graph_util.check_graph_equal(g1, g2) 18 | 19 | def test_super_resolution_example(): 20 | fname, symbol = super_resolution, super_resolution_sym 21 | compare_graph(fname, symbol, ishape=(1, 1, 224, 224)) 22 | 23 | if __name__ == '__main__': 24 | test_super_resolution_example() 25 | -------------------------------------------------------------------------------- /tests/python/unittest/test_symbol.py: -------------------------------------------------------------------------------- 1 | import nnvm.symbol as sym 2 | from nnvm import NNVMError 3 | 4 | def test_dense(): 5 | x = sym.Variable('x') 6 | y = sym.dense(x, units=30, name="fc") 7 | assert y.list_input_names() == ["x", "fc_weight", "fc_bias"] 8 | 9 | def test_batch_norm(): 10 | x = sym.Variable('x') 11 | y = sym.dense(x, units=30, name="fc") 12 | z = sym.batch_norm(x, name='bn') 13 | assert z.list_input_names('aux_state') == ['bn_moving_mean', 'bn_moving_var'] 14 | assert z.list_input_names('read_only') == ['x', 'bn_gamma', 'bn_beta'] 15 | 16 | def test_compose(): 17 | x = sym.Variable('x') 18 | z = sym.Variable('z') 19 | y = sym.exp(sym.elemwise_add(x, x, name='add', gpu=2), 20 | name='exp', gpu=1, attr={"kk": "1"}) 21 | 22 | assert y.list_input_names() == ['x'] 23 | assert y.list_output_names() == ["exp_output"] 24 | assert y.list_attr()['gpu'] == '1' 25 | z = y.get_internals() 26 | assert z['add_output'].list_output_names() == ['add_output'] 27 | assert y.list_attr(recursive=True)['add$gpu'] == '2' 28 | 29 | def test_default_input(): 30 | x = sym.Variable('x') 31 | y = sym.dense(data=x, units=30, name='fc', use_bias=False) 32 | assert y.list_input_names() == ['x', 'fc_weight'] 33 | tname = [z.list_output_names()[0] for z in y.list_input_variables()] 34 | assert tname == y.list_input_names() 35 | try: 36 | z = sym.elemwise_add(x) 37 | assert False 38 | except NNVMError: 39 | pass 40 | 41 | def test_copy(): 42 | x = sym.Variable('x') 43 | z = sym.Variable('z') 44 | y = sym.exp(sym.elemwise_add(x, x, name='add', gpu=2), 45 | name='exp', gpu=1, attr={"kk": "1"}) 46 | assert y.__copy__().debug_str() == y.debug_str() 47 | 48 | 49 | def test_op_name(): 50 | x = sym.Variable('x') 51 | y = sym.exp(x) 52 | op_name = y.attr("op_name") 53 | op_func = sym.__dict__[op_name] 54 | z = op_func(x) 55 | 56 | if __name__ == "__main__": 57 | test_op_name() 58 | test_copy() 59 | test_default_input() 60 | test_compose() 61 | test_batch_norm() 62 | -------------------------------------------------------------------------------- /tests/python/unittest/test_top_level1.py: -------------------------------------------------------------------------------- 1 | import nnvm.symbol as sym 2 | import nnvm.graph as graph 3 | 4 | def test_dense(): 5 | x = sym.Variable('x') 6 | x1 = sym.dense(x, units=3, name="dense") 7 | x2 = sym.flatten(x1) 8 | x3 = sym.softmax(x2) 9 | assert x3.list_input_names() == ['x', 'dense_weight', 'dense_bias'] 10 | 11 | 12 | def test_concatenate_split(): 13 | x = sym.Variable('x') 14 | y = sym.Variable('y') 15 | y = sym.concatenate(x, y) 16 | assert y.list_input_names() == ['x', 'y'] 17 | z = sym.split(y, indices_or_sections=10) 18 | assert len(z.list_output_names()) == 10 19 | z = sym.split(y, indices_or_sections=[10, 20]) 20 | assert len(z.list_output_names()) == 3 21 | 22 | def test_expand_dims(): 23 | x = sym.Variable('x') 24 | y = sym.expand_dims(x, axis=1, num_newaxis=2) 25 | assert y.list_input_names() == ['x'] 26 | 27 | 28 | def test_unary(): 29 | x = sym.Variable('x') 30 | x = sym.exp(x) 31 | x = sym.log(x) 32 | x = sym.sigmoid(x) 33 | x = sym.tanh(x) 34 | x = sym.relu(x) 35 | assert x.list_input_names() == ['x'] 36 | 37 | 38 | def test_batchnorm(): 39 | x = sym.Variable('x') 40 | x = sym.batch_norm(x, name="bn") 41 | assert x.list_input_names() == [ 42 | "x", "bn_gamma", "bn_beta", "bn_moving_mean", "bn_moving_var"] 43 | 44 | 45 | if __name__ == "__main__": 46 | test_concatenate_split() 47 | test_expand_dims() 48 | test_dense() 49 | test_unary() 50 | test_batchnorm() 51 | -------------------------------------------------------------------------------- /tests/python/unittest/test_top_level2.py: -------------------------------------------------------------------------------- 1 | import nnvm.symbol as sym 2 | 3 | def test_conv2d(): 4 | x = sym.Variable('x') 5 | y = sym.conv2d(x, channels=3, kernel_size=(3, 3), 6 | name="y", use_bias=False) 7 | assert y.list_input_names() == ["x", "y_weight"] 8 | 9 | 10 | def test_max_pool2d(): 11 | x = sym.Variable('x') 12 | y = sym.max_pool2d(x, pool_size=(3, 3), name="y") 13 | y = sym.global_max_pool2d(y) 14 | assert y.list_input_names() == ["x"] 15 | 16 | 17 | if __name__ == "__main__": 18 | test_conv2d() 19 | test_max_pool2d() 20 | -------------------------------------------------------------------------------- /tests/python/unittest/test_top_level3.py: -------------------------------------------------------------------------------- 1 | import nnvm.symbol as sym 2 | 3 | def test_reshape(): 4 | x = sym.Variable("x") 5 | y = sym.reshape(x, shape=(10, 20), name="y") 6 | assert(y.list_input_names() == ["x"]) 7 | 8 | 9 | def test_scalar_op(): 10 | x = sym.Variable("x") 11 | y = (1 / (x * 2) - 1) ** 2 12 | assert(y.list_input_names() == ["x"]) 13 | 14 | def test_leaky_relu(): 15 | x = sym.Variable("x") 16 | y = sym.leaky_relu(x, alpha=0.1) 17 | assert(y.list_input_names() == ["x"]) 18 | 19 | def test_prelu(): 20 | x = sym.Variable("x") 21 | w = sym.Variable("w") 22 | y = sym.prelu(x, w) 23 | assert(y.list_input_names()[0] == 'x') 24 | assert(y.list_input_names()[1] == 'w') 25 | 26 | if __name__ == "__main__": 27 | test_scalar_op() 28 | test_reshape() 29 | test_leaky_relu() 30 | test_prelu() 31 | -------------------------------------------------------------------------------- /tests/python/unittest/test_top_level4.py: -------------------------------------------------------------------------------- 1 | import nnvm.symbol as sym 2 | 3 | def test_binary_broadcast(): 4 | x = sym.Variable('x') 5 | y = sym.Variable('y') 6 | z = x + y 7 | z = x * y 8 | z = x - y 9 | z = x / y 10 | 11 | 12 | def test_broadcast_to(): 13 | x = sym.Variable('x') 14 | y = sym.broadcast_to(x, shape=(3, 3)) 15 | assert y.list_input_names() == ["x"] 16 | 17 | 18 | if __name__ == "__main__": 19 | test_binary_broadcast() 20 | test_broadcast_to() 21 | -------------------------------------------------------------------------------- /tests/scripts/task_build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | echo "Build TVM..." 3 | cd tvm 4 | cp make/config.mk . 5 | echo USE_CUDNN=1 >> config.mk 6 | echo USE_CUDA=1 >> config.mk 7 | echo USE_OPENCL=1 >> config.mk 8 | echo LLVM_CONFIG=llvm-config-4.0 >> config.mk 9 | echo USE_RPC=1 >> config.mk 10 | echo USE_BLAS=openblas >> config.mk 11 | echo USE_GRAPH_RUNTIME=1 >> config.mk 12 | make "$@" 13 | cd .. 14 | 15 | echo "Build NNVM..." 16 | make "$@" 17 | -------------------------------------------------------------------------------- /tests/scripts/task_clean.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | echo "Cleanup data..." 3 | cd tvm 4 | make clean 5 | cd .. 6 | make clean 7 | -------------------------------------------------------------------------------- /tests/scripts/task_frontend_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export PYTHONPATH=python:tvm/python:tvm/topi/python 4 | 5 | echo "Running ONNX frontend test..." 6 | python -m nose -v tests/python/frontend/onnx || exit -1 7 | 8 | echo "Running MXNet frontend test..." 9 | python -m nose -v tests/python/frontend/mxnet || exit -1 10 | 11 | echo "Running Keras frontend test..." 12 | python -m nose -v tests/python/frontend/keras || exit -1 13 | -------------------------------------------------------------------------------- /tests/scripts/task_lint.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | echo "Check codestyle of c++ code..." 3 | make cpplint || exit -1 4 | echo "Check codestyle of python code..." 5 | make pylint || exit -1 6 | echo "Check documentations of c++ code..." 7 | make doc 2>log.txt 8 | (cat log.txt| grep -v ENABLE_PREPROCESSING |grep -v "unsupported tag") > logclean.txt 9 | echo "---------Error Log----------" 10 | cat logclean.txt 11 | echo "----------------------------" 12 | (cat logclean.txt|grep warning) && exit -1 13 | (cat logclean.txt|grep error) && exit -1 14 | rm logclean.txt 15 | rm log.txt 16 | -------------------------------------------------------------------------------- /tests/scripts/task_python_docs.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | mkdir -p docs/_build/html 3 | # C++ doc 4 | make doc 5 | 6 | rm -rf python/nnvm/*.pyc python/nnvm/*/*.pyc 7 | 8 | cd docs 9 | PYTHONPATH=../python:../tvm/python:../tvm/topi/python make html || exit -1 10 | cd _build/html 11 | tar czf docs.tgz * 12 | mv docs.tgz ../../../ 13 | -------------------------------------------------------------------------------- /tests/scripts/task_python_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export PYTHONPATH=python:tvm/python:tvm/topi/python 4 | 5 | echo "Running unittest..." 6 | python -m nose -v tests/python/unittest || exit -1 7 | python3 -m nose -v tests/python/unittest || exit -1 8 | 9 | echo "Running compiler test..." 10 | python -m nose -v tests/python/compiler || exit -1 11 | python3 -m nose -v tests/python/compiler || exit -1 12 | -------------------------------------------------------------------------------- /tests/travis/run_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | if [ ${TASK} == "lint" ]; then 5 | make lint || exit -1 6 | echo "Check documentations of c++ code..." 7 | make doc 2>log.txt 8 | (cat log.txt| grep -v ENABLE_PREPROCESSING |grep -v "unsupported tag") > logclean.txt 9 | echo "---------Error Log----------" 10 | cat logclean.txt 11 | echo "----------------------------" 12 | (cat logclean.txt|grep warning) && exit -1 13 | (cat logclean.txt|grep error) && exit -1 14 | exit 0 15 | fi 16 | 17 | 18 | if [ ! ${TRAVIS_OS_NAME} == "osx" ]; then 19 | # use g++-4.8 for linux 20 | if [ ${CXX} == "g++" ]; then 21 | export CXX=g++-4.8 22 | fi 23 | fi 24 | 25 | if [ ${TASK} == "cpp_test" ]; then 26 | make -f dmlc-core/scripts/packages.mk gtest 27 | echo "GTEST_PATH="${CACHE_PREFIX} >> config.mk 28 | make test || exit -1 29 | for test in tests/cpp/*_test; do 30 | ./$test || exit -1 31 | done 32 | exit 0 33 | fi 34 | 35 | # run two test one for cython, one for ctypes 36 | if [ ${TASK} == "python_test" ]; then 37 | make clean 38 | make -j all || exit -1 39 | if [ ${TRAVIS_OS_NAME} == "osx" ]; then 40 | python -m nose tests/python/unittest/ || exit -1 41 | python3 -m nose tests/python/unittest/ || exit -1 42 | else 43 | nosetests tests/python/unittest/ || exit -1 44 | nosetests3 tests/python/unittest/ || exit -1 45 | fi 46 | 47 | make cython || exit -1 48 | make cython3 || exit -1 49 | 50 | if [ ${TRAVIS_OS_NAME} == "osx" ]; then 51 | python -m nose tests/python/unittest/ || exit -1 52 | python3 -m nose tests/python/unittest/ || exit -1 53 | else 54 | nosetests tests/python/unittest/ || exit -1 55 | nosetests3 tests/python/unittest/ || exit -1 56 | fi 57 | exit 0 58 | fi 59 | -------------------------------------------------------------------------------- /tests/travis/setup.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | if [ ${TRAVIS_OS_NAME} == "osx" ]; then 5 | brew update 6 | brew install python3 7 | if [ ${TASK} == "python_test" ]; then 8 | python -m pip install --user nose numpy cython 9 | python3 -m pip install --user nose numpy cython 10 | fi 11 | fi 12 | 13 | if [ ${TASK} == "lint" ]; then 14 | pip install --user cpplint 'pylint==1.4.4' 'astroid==1.3.6' 15 | fi 16 | -------------------------------------------------------------------------------- /tests/travis/travis_after_failure.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | -------------------------------------------------------------------------------- /tutorials/README.txt: -------------------------------------------------------------------------------- 1 | Tutorials 2 | ========= 3 | This page contains the tutorials about NNVM. 4 | -------------------------------------------------------------------------------- /tutorials/from_coreml.py: -------------------------------------------------------------------------------- 1 | """ 2 | Compile CoreML Models 3 | ===================== 4 | **Author**: `Joshua Z. Zhang `_ 5 | 6 | This article is an introductory tutorial to deploy CoreML models with NNVM. 7 | 8 | For us to begin with, coremltools module is required to be installed. 9 | 10 | A quick solution is to install via pip 11 | ```bash 12 | pip install -U coremltools --user 13 | ``` 14 | or please refer to official site 15 | https://github.com/apple/coremltools 16 | """ 17 | import nnvm 18 | import tvm 19 | import coremltools as cm 20 | import numpy as np 21 | from PIL import Image 22 | 23 | def download(url, path, overwrite=False): 24 | import os 25 | if os.path.isfile(path) and not overwrite: 26 | print('File {} existed, skip.'.format(path)) 27 | return 28 | print('Downloading from url {} to {}'.format(url, path)) 29 | try: 30 | import urllib.request 31 | urllib.request.urlretrieve(url, path) 32 | except: 33 | import urllib 34 | urllib.urlretrieve(url, path) 35 | 36 | ###################################################################### 37 | # Load pretrained CoreML model 38 | # ---------------------------- 39 | # We will download and load a pretrained mobilenet classification network 40 | # provided by apple in this example 41 | model_url = 'https://docs-assets.developer.apple.com/coreml/models/MobileNet.mlmodel' 42 | model_file = 'mobilenet.mlmodel' 43 | download(model_url, model_file) 44 | # now you mobilenet.mlmodel on disk 45 | mlmodel = cm.models.MLModel(model_file) 46 | # we can load the graph as NNVM compatible model 47 | sym, params = nnvm.frontend.from_coreml(mlmodel) 48 | 49 | ###################################################################### 50 | # Load a test image 51 | # ------------------ 52 | # A single cat dominates the examples! 53 | from PIL import Image 54 | img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true' 55 | download(img_url, 'cat.png') 56 | img = Image.open('cat.png').resize((224, 224)) 57 | #x = np.transpose(img, (2, 0, 1))[np.newaxis, :] 58 | image = np.asarray(img) 59 | image = image.transpose((2, 0, 1)) 60 | x = image[np.newaxis, :] 61 | ###################################################################### 62 | # Compile the model on NNVM 63 | # --------------------------- 64 | # We should be familiar with the process right now. 65 | import nnvm.compiler 66 | target = 'cuda' 67 | shape_dict = {'image': x.shape} 68 | graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, params=params) 69 | 70 | ###################################################################### 71 | # Execute on TVM 72 | # ------------------- 73 | # The process is no different from other example 74 | from tvm.contrib import graph_runtime 75 | ctx = tvm.gpu(0) 76 | dtype = 'float32' 77 | m = graph_runtime.create(graph, lib, ctx) 78 | # set inputs 79 | m.set_input('image', tvm.nd.array(x.astype(dtype))) 80 | m.set_input(**params) 81 | # execute 82 | m.run() 83 | # get outputs 84 | output_shape = (1000,) 85 | tvm_output = m.get_output(0, tvm.nd.empty(output_shape, dtype)).asnumpy() 86 | top1 = np.argmax(tvm_output) 87 | 88 | ##################################################################### 89 | # Look up synset name 90 | # ------------------- 91 | # Look up prdiction top 1 index in 1000 class synset. 92 | synset_url = ''.join(['https://gist.githubusercontent.com/zhreshold/', 93 | '4d0b62f3d01426887599d4f7ede23ee5/raw/', 94 | '596b27d23537e5a1b5751d2b0481ef172f58b539/', 95 | 'imagenet1000_clsid_to_human.txt']) 96 | synset_name = 'synset.txt' 97 | download(synset_url, synset_name) 98 | with open(synset_name) as f: 99 | synset = eval(f.read()) 100 | print('Top-1 id', top1, 'class name', synset[top1]) 101 | -------------------------------------------------------------------------------- /tutorials/imagenet_inference_gpu.py: -------------------------------------------------------------------------------- 1 | """ 2 | Compile GPU Inference 3 | ===================== 4 | **Author**: `Yuwei Hu `_ 5 | 6 | This is an example of using NNVM to compile MobileNet/ResNet model and deploy its inference on GPU. 7 | 8 | To begin with, we import nnvm(for compilation) and TVM(for deployment). 9 | """ 10 | import tvm 11 | import numpy as np 12 | from tvm.contrib import nvcc, graph_runtime 13 | import nnvm.compiler 14 | import nnvm.testing 15 | 16 | ###################################################################### 17 | # Register the NVCC Compiler Option 18 | # --------------------------------- 19 | # NNVM optimizes the graph and relies on TVM to generate fast GPU code. 20 | # To get the maximum performance, we need to enable nvcc's compiler hook. 21 | # This usually gives better performance than nvrtc mode. 22 | 23 | @tvm.register_func 24 | def tvm_callback_cuda_compile(code): 25 | ptx = nvcc.compile_cuda(code, target="ptx") 26 | return ptx 27 | 28 | ###################################################################### 29 | # Prepare the Benchmark 30 | # --------------------- 31 | # We construct a standard imagenet inference benchmark. 32 | # NNVM needs two things to compile a deep learning model: 33 | # 34 | # - net: the graph representation of the computation 35 | # - params: a dictionary of str to parameters 36 | # 37 | # We use nnvm's testing utility to produce the model description and random parameters 38 | # so that the example does not depend on a specific front-end framework. 39 | # 40 | # .. note:: 41 | # 42 | # In a typical workflow, we can get this pair from :any:`nnvm.frontend` 43 | # 44 | target = "cuda" 45 | ctx = tvm.gpu(0) 46 | batch_size = 1 47 | num_classes = 1000 48 | image_shape = (3, 224, 224) 49 | data_shape = (batch_size,) + image_shape 50 | out_shape = (batch_size, num_classes) 51 | # To use ResNet to do inference, run the following instead 52 | #net, params = nnvm.testing.resnet.get_workload( 53 | # batch_size=1, image_shape=image_shape) 54 | net, params = nnvm.testing.mobilenet.get_workload( 55 | batch_size=1, image_shape=image_shape) 56 | 57 | ###################################################################### 58 | # Compile the Graph 59 | # ----------------- 60 | # To compile the graph, we call the build function with the graph 61 | # configuration and parameters. 62 | # When parameters are provided, NNVM will pre-compute certain part of the graph if possible (e.g. simplify batch normalization to scale shift), 63 | # and return the updated parameters. 64 | 65 | graph, lib, params = nnvm.compiler.build( 66 | net, target, shape={"data": data_shape}, params=params) 67 | 68 | 69 | ###################################################################### 70 | # Run the Compiled Module 71 | # ----------------------- 72 | # 73 | # To deploy the module, we call :any:`tvm.contrib.graph_runtime.create` passing in the graph, the lib, and context. 74 | # Thanks to TVM, we can deploy the compiled module to many platforms and languages. 75 | # The deployment module is designed to contain minimum dependencies. 76 | # This example runs on the same machine. 77 | # 78 | # Note that the code below no longer depends on NNVM, and only relies TVM's runtime to run(deploy). 79 | data = np.random.uniform(-1, 1, size=data_shape).astype("float32") 80 | module = graph_runtime.create(graph, lib, ctx) 81 | # set input 82 | module.set_input(**params) 83 | module.set_input("data", data) 84 | # run 85 | module.run() 86 | # get output 87 | out = module.get_output(0, tvm.nd.empty(out_shape)) 88 | # convert to numpy 89 | out.asnumpy() 90 | --------------------------------------------------------------------------------