[TVM] Relax API (#576)

Добавлена поддержка Relax API для Apache TVM.

Author: Vladislav-Denisov

demo/benchmark_configs/TVM.xml

@@ -30,7 +30,7 @@
             <ChannelSwap></ChannelSwap>
             <Target>llvm</Target>
             <Layout>NCHW</Layout>
-            <VirtualMachine>False</VirtualMachine>
+            <HighLevelAPI>Relay</HighLevelAPI>
             <OptimizationLevel>3</OptimizationLevel>
         </FrameworkDependent>
     </Test>

docker/TVM/Dockerfile

@@ -5,37 +5,60 @@ WORKDIR /root/
 # Installing miniconda
 RUN wget -q --no-check-certificate -c https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     bash Miniconda3-latest-Linux-x86_64.sh -b && \
-    ./miniconda3/bin/conda create -n tvm-env -y python=3.7.16 && \
+    ./miniconda3/bin/conda create -n tvm-env -y python=3.8.20 && \
     rm -rf /root/miniconda3/pkgs/* && \
     rm ~/Miniconda3-latest-Linux-x86_64.sh -f
-    
+
 ENV PATH /root/miniconda3/envs/tvm-env/bin:/root/miniconda3/bin:$PATH
 RUN echo "source activate tvm-env" > ~/.bashrc
 RUN export LD_LIBRARY_PATH=/root/miniconda3/envs/tvm-env/lib:${LD_LIBRARY_PATH}
 RUN conda config --add channels intel
 
 # Installing dependencies
-RUN python3 -m pip install pycocotools docker PyYAML
+RUN python3 -m pip install pycocotools docker PyYAML gluoncv[full] opencv-python Cython psutil
 RUN apt-get update && apt-get install -y ffmpeg libsm6 libxext6
-RUN python3 -m pip install gluoncv[full]
-RUN python3 -m pip install opencv-python
-   
+RUN apt-get update && apt-get install -y -qq --no-install-recommends cmake software-properties-common
+RUN apt-add-repository "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-19 main" && \
+    wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
+RUN apt-get update && \
+    apt-get install -y -qq --no-install-recommends llvm-19 llvm-19-dev && \
+    rm -rf /var/lib/apt/lists/*
+
 # Installing Apache-TVM 
-RUN python3 -m pip install apache-tvm==0.14.dev264
+ARG TVM_VERSION=v0.20.0
+
+RUN git clone --recursive https://github.com/apache/tvm tvm --branch ${TVM_VERSION} --single-branch
+ENV TVM_BUILD_DIR=/tmp/build-tvm
+RUN mkdir $TVM_BUILD_DIR && cd $TVM_BUILD_DIR
+RUN cp /root/tvm/cmake/config.cmake .
+RUN echo "set(CMAKE_BUILD_TYPE RelWithDebInfo)" >> config.cmake
+RUN echo "set(USE_LLVM /usr/bin/llvm-config-19)" >> config.cmake
+RUN echo "set(HIDE_PRIVATE_SYMBOLS ON)" >> config.cmake
+RUN echo "set(USE_CUDA    OFF)" >> config.cmake
+RUN echo "set(USE_METAL   OFF)" >> config.cmake
+RUN echo "set(USE_VULKAN  OFF)" >> config.cmake
+RUN echo "set(USE_OPENCL  OFF)" >> config.cmake
+RUN echo "set(USE_CUBLAS  OFF)" >> config.cmake
+RUN echo "set(USE_CUDNN   OFF)" >> config.cmake
+RUN echo "set(USE_CUTLASS OFF)" >> config.cmake
+
+RUN /bin/bash -c 'cmake /root/tvm/ && cmake --build . -- -j$(nproc --all)'
+
+WORKDIR /root/
+RUN export TVM_LIBRARY_PATH=$TVM_BUILD_DIR
+RUN python3 -m pip install -e /root/tvm/python
 
 # ARG
 ARG TORCH_VERSION=2.0.1
 ARG TORCHVISION_VERSION=0.15.2
 ARG TFLite_VERSION=2.14.0
 ARG MXNET_VERSION=1.9.1
 ARG ONNX_VERSION=1.15.0
-ARG CAFFE_VERSION=1.1.0
 
 ARG PyTorch
 ARG TFLite
 ARG MXNet
 ARG ONNX
-ARG CAFFE
 
 RUN if [ "${PyTorch}" = "true"  ]; then \
     python3 -m pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION}; \
@@ -65,14 +88,6 @@ RUN if [ "$ONNX" = "true"  ]; then \
     python3 -m pip install onnxruntime==${ONNX_VERSION} && python3 -m pip install onnx==${ONNX_VERSION}; \
     fi
 
-RUN if [ "$CAFFE" = "true"  ]; then \
-    /bin/bash -c "conda install -n tvm-env -y -c intel openvino-ie4py-ubuntu20=2022.1.0"; \
-    /bin/bash -c "conda install -n tvm-env -y -c intel scikit-image=0.17.2"; \
-    /bin/bash -c "conda install -n tvm-env -y -c intel caffe"; \
-    /bin/bash -c "conda install -n tvm-env -y -c anaconda libgcc-ng=11.2.0"; \
-    /bin/bash -c "conda install -n tvm-env -y -c defaults protobuf=3.17.2 libprotobuf=3.17.2 requests=2.27.1"; \
-    fi
-
 WORKDIR /tmp/open_model_zoo/
 RUN git remote add omz_custom_tvm https://github.com/itlab-vision/open_model_zoo_tvm.git && \ 
     git fetch omz_custom_tvm && \

src/benchmark/frameworks/tvm/tvm_parameters_parser.py

@@ -15,7 +15,7 @@ def parse_parameters(self, curr_test):
         CONFIG_FRAMEWORK_DEPENDENT_STD_TAG = 'Std'
         CONFIG_FRAMEWORK_DEPENDENT_CHANNEL_SWAP_TAG = 'ChannelSwap'
         CONFIG_FRAMEWORK_DEPENDENT_LAYOUT_TAG = 'Layout'
-        CONFIG_FRAMEWORK_DEPENDENT_VIRTUAL_MACHINE = 'VirtualMachine'
+        CONFIG_FRAMEWORK_DEPENDENT_HIGH_LEVEL_API = 'HighLevelAPI'
 
         dep_parameters_tag = curr_test.getElementsByTagName(CONFIG_FRAMEWORK_DEPENDENT_TAG)[0]
 
@@ -39,8 +39,8 @@ def parse_parameters(self, curr_test):
             CONFIG_FRAMEWORK_DEPENDENT_LAYOUT_TAG)[0].firstChild
         _target = dep_parameters_tag.getElementsByTagName(
             CONFIG_FRAMEWORK_DEPENDENT_TARGET)[0].firstChild
-        _vm = dep_parameters_tag.getElementsByTagName(
-            CONFIG_FRAMEWORK_DEPENDENT_VIRTUAL_MACHINE)[0].firstChild
+        _high_level_api = dep_parameters_tag.getElementsByTagName(
+            CONFIG_FRAMEWORK_DEPENDENT_HIGH_LEVEL_API)[0].firstChild
 
         return TVMParameters(
             framework=_framework.data if _framework else None,
@@ -53,14 +53,15 @@ def parse_parameters(self, curr_test):
             optimization_level=_optimization_level.data if _optimization_level else None,
             layout=_layout.data if _layout else None,
             target=_target.data if _target else None,
-            vm=_vm.data if _vm else None,
+            high_level_api=_high_level_api.data if _high_level_api else None,
         )
 
 
 class TVMParameters(FrameworkParameters):
     def __init__(self, framework, input_name, input_shape,
                  normalize, mean, std, channel_swap,
-                 optimization_level, layout, target, vm):
+                 optimization_level, layout, target,
+                 high_level_api):
         self.framework = None
         self.input_name = None
         self.input_shape = None
@@ -71,7 +72,7 @@ def __init__(self, framework, input_name, input_shape,
         self.optimization_level = None
         self.layout = None
         self.target = 'llvm'
-        self.vm = None
+        self.high_level_api = None
 
         if self._framework_is_correct(framework):
             self.framework = framework
@@ -93,8 +94,8 @@ def __init__(self, framework, input_name, input_shape,
             self.layout = layout
         if self._parameter_is_not_none(target):
             self.target = target
-        if self._parameter_is_not_none(vm):
-            self.vm = vm
+        if self._parameter_is_not_none(high_level_api):
+            self.high_level_api = high_level_api
 
     @staticmethod
     def _framework_is_correct(framework):

src/benchmark/frameworks/tvm/tvm_process.py

@@ -57,11 +57,6 @@ def _fill_command_line(self):
             common_params = TVMProcess._add_flag_to_cmd_line(
                 common_params, '--norm')
 
-        vm = self._test.dep_parameters.vm
-        if vm == 'True':
-            common_params = TVMProcess._add_flag_to_cmd_line(
-                common_params, '-vm')
-
         mean = self._test.dep_parameters.mean
         common_params = TVMProcess._add_optional_argument_to_cmd_line(
             common_params, '--mean', mean)
@@ -86,6 +81,10 @@ def _fill_command_line(self):
         common_params = TVMProcess._add_optional_argument_to_cmd_line(
             common_params, '--target', target)
 
+        high_level_api = self._test.dep_parameters.high_level_api
+        common_params = TVMProcess._add_optional_argument_to_cmd_line(
+            common_params, '--high_level_api', high_level_api)
+
         return f'{common_params}'
 
 

src/configs/README.md

@@ -354,6 +354,7 @@
     изображения.
   - `ChannelSwap` - тег, необязательный для заполнения. Описывает изменение порядка каналов на
     входном изображении. По умолчанию будет установлен порядок (2, 0, 1), что соответствует BGR.
+  - `HighLevelAPI` - тег, необязательный для заполнения. Определяет используемое высокоуровневое API: `Relay`, `RelayVM` или `RelaxVM`. По умолчанию задается значение `Relay`.
   - `OptimizationLevel` - тег, необязательный для заполнения. Определяет уровень оптимизаций для
     графа вычислений, которые ускоряют инференс. По умолчанию оптимизации не применяются.
   - `Framework` - тег, обязательный для заполнения. Определяет фреймворк, модели которого будут
@@ -766,7 +767,7 @@
         <ChannelSwap></ChannelSwap>
         <Layout>NCHW</Layout>
         <Target>llvm</Target>
-        <VirtualMachine>True</VirtualMachine>
+        <HighLevelAPI>RelaxVM</HighLevelAPI>
         <OptimizationLevel>3</OptimizationLevel>
     </FrameworkDependent>
 </Test>

src/configs/benchmark_configuration_file_template.xml

@@ -446,41 +446,7 @@
             <Std></Std>
             <Layout></Layout>
             <ChannelSwap></ChannelSwap>
-            <VirtualMachine></VirtualMachine>
-            <Target></Target>
-            <OptimizationLevel></OptimizationLevel>
-        </FrameworkDependent>
-    </Test>
-    <Test>
-        <Model>
-            <Task></Task>
-            <Name></Name>
-            <Precision></Precision>
-            <SourceFramework>TVM</SourceFramework>
-            <ModelPath></ModelPath>
-            <WeightsPath></WeightsPath>
-        </Model>
-        <Dataset>
-            <Name></Name>
-            <Path></Path>
-        </Dataset>
-        <FrameworkIndependent>
-            <InferenceFramework>TVM</InferenceFramework>
-            <BatchSize></BatchSize>
-            <Device></Device>
-            <IterationCount></IterationCount>
-            <TestTimeLimit></TestTimeLimit>
-        </FrameworkIndependent>
-        <FrameworkDependent>
-            <InputName></InputName>
-            <InputShape></InputShape>
-            <Framework></Framework> <!--Принимает значения: caffe, onnx, pytorch, mxnet, tvm, tflite-->
-            <Normalize></Normalize>
-            <Mean></Mean>
-            <Std></Std>
-            <Layout></Layout>
-            <ChannelSwap></ChannelSwap>
-            <VirtualMachine></VirtualMachine>
+            <HighLevelAPI></HighLevelAPI>
             <Target></Target>
             <OptimizationLevel></OptimizationLevel>
         </FrameworkDependent>

src/inference/README.md

@@ -956,6 +956,7 @@ inference_tvm.py
   ImageNet.
 - `-ni / --number_iter` - количество прямых проходов по сети.
   По умолчанию выполняется один проход по сети.
+- `--high_level_api` - высокоуровневое API: `Relay`, `RelayVM` или `RelaxVM`. По умолчанию `Relay`.
 - `-ol / --opt_level` - параметр, определяющий уровень оптимизации
   графа вычислений нейронной сети для ускорения инференса. По умолчанию
   оптимизации не применяются.

src/inference/inference_tvm.py

@@ -166,10 +166,12 @@ def cli_argument_parser():
                         default=0.5,
                         type=float,
                         dest='threshold')
-    parser.add_argument('-vm', '--virtual_machine',
-                        help='Flag to use VirtualMachine API',
-                        action='store_true',
-                        dest='vm')
+    parser.add_argument('--high_level_api',
+                        help='Type of high level API',
+                        choices=['Relay', 'RelayVM', 'RelaxVM'],
+                        default='Relay',
+                        type=str,
+                        dest='high_level_api')
     parser.add_argument('--raw_output',
                         help='Raw output without logs.',
                         default=False,
@@ -209,8 +211,7 @@ def main():
                                            args.input_name,
                                            io.get_slice_input,
                                            args.time,
-                                           args.vm)
-
+                                           args.high_level_api)
         if not args.raw_output:
             if args.number_iter == 1:
                 try:

src/inference/tvm_auxiliary.py

@@ -22,11 +22,13 @@ def _infer_slice(self, input_name, module, slice_input):
         pass
 
     @staticmethod
-    def get_helper(vm):
-        if vm:
-            return InferenceVMApi()
-        else:
+    def get_helper(high_level_api):
+        if high_level_api == 'Relay':
             return InferenceRelayAPI()
+        elif high_level_api == 'RelayVM':
+            return InferenceRelayVMApi()
+        elif high_level_api == 'RelaxVM':
+            return InferenceRelaxVMApi()
 
     @abc.abstractmethod
     def _inference_tvm(self, module, input_name, slice_input):
@@ -77,22 +79,37 @@ def _inference_tvm(self, module, input_name, slice_input):
         return [module.get_output(i) for i in range(num_of_outputs)]
 
 
-class InferenceVMApi(InferenceHelper):
+class InferenceRelayVMApi(InferenceHelper):
     def __init__(self):
         super().__init__()
 
     def _infer_slice(self, input_name, module, slice_input):
         module.set_input('main', slice_input[input_name])
         module.run()
-        res = module.get_outputs()
-        return res
+        return module.get_outputs()
 
     def _inference_tvm(self, module, input_name, slice_input):
         module.set_input('main', slice_input[input_name])
         module.run()
         return module.get_outputs()
 
 
+class InferenceRelaxVMApi(InferenceHelper):
+    def __init__(self):
+        super().__init__()
+
+    def _infer_slice(self, input_name, module, slice_input):
+        module.set_input('main', slice_input[input_name])
+        module.invoke_stateful('main')
+        return module.get_outputs('main')
+
+    def _inference_tvm(self, module, input_name, slice_input):
+        import tvm
+        module.set_input('main', slice_input[input_name])
+        module.invoke_stateful('main')
+        return [tvm.nd.array(module.get_outputs('main'))]
+
+
 class OutputPreparer:
     def __init__(self, source_framework):
         self.source_framework = source_framework
@@ -163,7 +180,7 @@ def create_dict_for_converter(args):
         'opt_level': args.opt_level,
         'target': args.target,
         'module': args.module,
-        'vm': args.vm,
+        'high_level_api': args.high_level_api,
         'source_framework': args.source_framework,
     }
     return dictionary
@@ -199,8 +216,8 @@ def create_dict_for_output_preparer(args):
     return dictionary
 
 
-def inference_tvm(module, num_of_iterations, input_name, get_slice, test_duration, vm):
-    inference_helper = InferenceHelper.get_helper(vm)
+def inference_tvm(module, num_of_iterations, input_name, get_slice, test_duration, high_level_api):
+    inference_helper = InferenceHelper.get_helper(high_level_api)
     return inference_helper.inference_tvm(module, num_of_iterations, input_name, get_slice, test_duration)
 
 

src/model_converters/tvm_converter/README.md

@@ -68,7 +68,7 @@ tvm_compiler.py --mod <model> \
                 --params <parameters> \
                 --target <target> \
                 --opt_level <opt_level> \
-                --virtual_machine <virtual_machine> \
+                --high_level_api <high_level_api> \
                 --lib_name <lib_name> \
                 --output_dir <output_dir>
 ```
@@ -82,7 +82,7 @@ for the Relay API or to the `.so`+`.ro` format for the VirtualMachine API.
 - `-p / --params` is a path to an `.params` file with a model parameters.
 - `-t / --target` is target device information, for example `llvm` for CPU.
 - `--opt_level` is the optimization level of the task extractions.
-- `-vm / --virtual_machine` is a flag to use the VirtualMachine API.
+- `--high_level_api` is a high level API: `Relay`, `RelayVM`, `RelaxVM`.
 - `--lib_name` is a file name to save compiled model.
 - `-op / --output_dir` is a path to save the model.
 

src/model_converters/tvm_converter/tvm_auxiliary/caffe_format.py

@@ -23,9 +23,11 @@ def _convert_model_from_framework(self):
 
         with open(self.model_params, 'rb') as f:
             init_net.ParseFromString(f.read())
-
-        model, params = self.tvm.relay.frontend.from_caffe(init_net,
-                                                           predict_net,
-                                                           shape_dict,
-                                                           dtype_dict)
-        return model, params
+        if self.high_level_api in ['Relay', 'RelayVM']:
+            model, params = self.tvm.relay.frontend.from_caffe(init_net,
+                                                               predict_net,
+                                                               shape_dict,
+                                                               dtype_dict)
+            return model, params
+        else:
+            raise ValueError(f'API {self.high_level_api} is not supported')