From 8cc969399a48d6379c85179279ff131bee15dd23 Mon Sep 17 00:00:00 2001 From: grimoire Date: Fri, 28 Jan 2022 18:35:50 +0800 Subject: [PATCH 1/2] add passes --- csrc/backend_ops/torchscript/CMakeLists.txt | 41 +-------- .../torchscript/ops/CMakeLists.txt | 41 +++++++++ .../modulated_deform_conv_cpu.cpp | 0 .../modulated_deform_conv_cuda.cu | 0 .../torchscript/optimizer/CMakeLists.txt | 12 +++ .../torchscript/optimizer/main.cpp | 86 +++++++++++++++++++ .../torchscript/optimizer/optimizer.cpp | 31 +++++++ .../torchscript/optimizer/optimizer.h | 7 ++ 8 files changed, 179 insertions(+), 39 deletions(-) create mode 100644 csrc/backend_ops/torchscript/ops/CMakeLists.txt rename csrc/backend_ops/torchscript/{ => ops}/modulated_deform_conv/modulated_deform_conv_cpu.cpp (100%) rename csrc/backend_ops/torchscript/{ => ops}/modulated_deform_conv/modulated_deform_conv_cuda.cu (100%) create mode 100644 csrc/backend_ops/torchscript/optimizer/CMakeLists.txt create mode 100644 csrc/backend_ops/torchscript/optimizer/main.cpp create mode 100644 csrc/backend_ops/torchscript/optimizer/optimizer.cpp create mode 100644 csrc/backend_ops/torchscript/optimizer/optimizer.h diff --git a/csrc/backend_ops/torchscript/CMakeLists.txt b/csrc/backend_ops/torchscript/CMakeLists.txt index 36c26845bb..e383129992 100644 --- a/csrc/backend_ops/torchscript/CMakeLists.txt +++ b/csrc/backend_ops/torchscript/CMakeLists.txt @@ -1,42 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. cmake_minimum_required(VERSION 3.14) -include(${CMAKE_SOURCE_DIR}/cmake/cuda.cmake NO_POLICY_SCOPE) -if("cuda" IN_LIST MMDEPLOY_TARGET_DEVICES) - project(mmdeploy_torchscript_ops CUDA CXX) - include(${CMAKE_SOURCE_DIR}/cmake/cuda.cmake NO_POLICY_SCOPE) - file(GLOB_RECURSE BACKEND_OPS_SRCS *.cpp *.cu) -else() - project(mmdeploy_torchscript_ops CXX) - file(GLOB_RECURSE BACKEND_OPS_SRCS *.cpp) -endif() - -include(${CMAKE_SOURCE_DIR}/cmake/common.cmake) -find_package(Torch REQUIRED) - -set_targets(${PROJECT_NAME} BACKEND_OPS_OBJ BACKEND_OPS_STATIC BACKEND_OPS_MODULE) - -build_object_target(${BACKEND_OPS_OBJ} "${BACKEND_OPS_SRCS}") -target_compile_definitions(${BACKEND_OPS_OBJ} - PRIVATE -DTHRUST_IGNORE_DEPRECATED_CPP_DIALECT=1) -target_include_directories(${BACKEND_OPS_OBJ} - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../common) -target_include_directories(${BACKEND_OPS_OBJ} - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/common) - -if("cuda" IN_LIST MMDEPLOY_TARGET_DEVICES) -target_include_directories(${BACKEND_OPS_OBJ} - PRIVATE ${CUDA_TOOLKIT_ROOT_DIR}/include) -endif() -target_link_libraries(${BACKEND_OPS_OBJ} PRIVATE ${TORCH_LIBRARIES}) - -# Build module library. It is used to inference with torchscript -build_module_target(${BACKEND_OPS_MODULE} ${BACKEND_OPS_OBJ} "PRIVATE") -add_library(mmdeploy::torchscript_ops ALIAS ${BACKEND_OPS_MODULE}) -install_targets(${BACKEND_OPS_MODULE}) - -if (MMDEPLOY_BUILD_SDK) - ## Build static library. SDK's uses it to build `trt_net` module - build_static_target(${BACKEND_OPS_STATIC} ${BACKEND_OPS_OBJ} "PRIVATE") - add_library(mmdeploy::torchscript_ops::static ALIAS ${BACKEND_OPS_STATIC}) -endif () +add_subdirectory(ops) +add_subdirectory(optimizer) diff --git a/csrc/backend_ops/torchscript/ops/CMakeLists.txt b/csrc/backend_ops/torchscript/ops/CMakeLists.txt new file mode 100644 index 0000000000..f9d54d31b1 --- /dev/null +++ b/csrc/backend_ops/torchscript/ops/CMakeLists.txt @@ -0,0 +1,41 @@ +# Copyright (c) OpenMMLab. All rights reserved. +cmake_minimum_required(VERSION 3.14) + +if("cuda" IN_LIST MMDEPLOY_TARGET_DEVICES) + project(mmdeploy_torchscript_ops CUDA CXX) + include(${CMAKE_SOURCE_DIR}/cmake/cuda.cmake NO_POLICY_SCOPE) + file(GLOB_RECURSE BACKEND_OPS_SRCS *.cpp *.cu) +else() + project(mmdeploy_torchscript_ops CXX) + file(GLOB_RECURSE BACKEND_OPS_SRCS *.cpp) +endif() + +include(${CMAKE_SOURCE_DIR}/cmake/common.cmake) +find_package(Torch REQUIRED) + +set_targets(${PROJECT_NAME} BACKEND_OPS_OBJ BACKEND_OPS_STATIC BACKEND_OPS_MODULE) + +build_object_target(${BACKEND_OPS_OBJ} "${BACKEND_OPS_SRCS}") +target_compile_definitions(${BACKEND_OPS_OBJ} + PRIVATE -DTHRUST_IGNORE_DEPRECATED_CPP_DIALECT=1) +target_include_directories(${BACKEND_OPS_OBJ} + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../common) +target_include_directories(${BACKEND_OPS_OBJ} + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/common) + +if("cuda" IN_LIST MMDEPLOY_TARGET_DEVICES) + target_include_directories(${BACKEND_OPS_OBJ} + PRIVATE ${CUDA_TOOLKIT_ROOT_DIR}/include) +endif() +target_link_libraries(${BACKEND_OPS_OBJ} PRIVATE ${TORCH_LIBRARIES}) + +# Build module library. It is used to inference with torchscript +build_module_target(${BACKEND_OPS_MODULE} ${BACKEND_OPS_OBJ} "PRIVATE") +add_library(mmdeploy::torchscript_ops ALIAS ${BACKEND_OPS_MODULE}) +install_targets(${BACKEND_OPS_MODULE}) + +if (MMDEPLOY_BUILD_SDK) + ## Build static library. + build_static_target(${BACKEND_OPS_STATIC} ${BACKEND_OPS_OBJ} "PRIVATE") + add_library(mmdeploy::torchscript_ops::static ALIAS ${BACKEND_OPS_STATIC}) +endif () diff --git a/csrc/backend_ops/torchscript/modulated_deform_conv/modulated_deform_conv_cpu.cpp b/csrc/backend_ops/torchscript/ops/modulated_deform_conv/modulated_deform_conv_cpu.cpp similarity index 100% rename from csrc/backend_ops/torchscript/modulated_deform_conv/modulated_deform_conv_cpu.cpp rename to csrc/backend_ops/torchscript/ops/modulated_deform_conv/modulated_deform_conv_cpu.cpp diff --git a/csrc/backend_ops/torchscript/modulated_deform_conv/modulated_deform_conv_cuda.cu b/csrc/backend_ops/torchscript/ops/modulated_deform_conv/modulated_deform_conv_cuda.cu similarity index 100% rename from csrc/backend_ops/torchscript/modulated_deform_conv/modulated_deform_conv_cuda.cu rename to csrc/backend_ops/torchscript/ops/modulated_deform_conv/modulated_deform_conv_cuda.cu diff --git a/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt b/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt new file mode 100644 index 0000000000..3a0121bc82 --- /dev/null +++ b/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt @@ -0,0 +1,12 @@ +# Copyright (c) OpenMMLab. All rights reserved. +cmake_minimum_required(VERSION 3.14) +project(ts_optimizer) + +find_package(Torch REQUIRED) + +file(GLOB_RECURSE OPTIMIZER_SRCS *.cpp) + +add_executable(${PROJECT_NAME} ${OPTIMIZER_SRCS}) +target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES}) +target_link_directories(${PROJECT_NAME} PRIVATE mmdeploy::torchscript_ops) + diff --git a/csrc/backend_ops/torchscript/optimizer/main.cpp b/csrc/backend_ops/torchscript/optimizer/main.cpp new file mode 100644 index 0000000000..cfd69d025c --- /dev/null +++ b/csrc/backend_ops/torchscript/optimizer/main.cpp @@ -0,0 +1,86 @@ +#include +#include + +#include +#include + +#include "optimizer.h" + +typedef std::unordered_map ArgMap; + +std::string get_or_default(const ArgMap& args_map, const std::string& key, + const std::string& default_val) { + auto iter = args_map.find(key); + return iter != args_map.end() ? iter->second : default_val; +} + +static void help() { + fprintf(stderr, "Usage: ts_opt [-backend=backend_name] [-out=out_file] model_file\n"); +} + +static ArgMap parse_args(int argc, char* argv[]) { + ArgMap args_map; + std::string model_file_key = "__model_file__"; + + for (int arg_id = 1; arg_id < argc; ++arg_id) { + std::string arg_str(argv[arg_id]); + size_t pos_equ = arg_str.find('='); + std::string key; + if (pos_equ != std::string::npos) { + key = arg_str.substr(0, pos_equ); + } else { + pos_equ = -1; + key = model_file_key; + } + + if (args_map.count(key)) { + fprintf(stderr, "ERROR: duplicate key: %s\n", key.c_str()); + help(); + exit(-1); + } + + args_map[key] = arg_str.substr(pos_equ + 1); + } + + if (args_map.count(model_file_key) == 0) { + fprintf(stderr, "ERROR: model file is required."); + help(); + exit(-1); + } + + return args_map; +} + +int main(int argc, char* argv[]) { + if (argc < 2) { + help(); + return -1; + } + + auto args_map = parse_args(argc, argv); + + std::string backend = get_or_default(args_map, "-backend", "torchscript"); + std::string model_file = args_map["__model_file__"]; + std::string output_file = get_or_default(args_map, "-out", model_file); + + // TODO: Dynamic link custom extension + + torch::jit::script::Module model; + try { + model = torch::jit::load(model_file); + } catch (const c10::Error& e) { + fprintf(stderr, "ERROR: fail to load model from %s.\n", model_file.c_str()); + exit(-1); + } + + if (backend == "torchscript") { + model = mmdeploy::optimize_for_torchscript(model); + } else { + fprintf(stderr, "No optimize for backend: %s\n", backend.c_str()); + exit(-1); + } + + model.save(output_file); + + return 0; +} \ No newline at end of file diff --git a/csrc/backend_ops/torchscript/optimizer/optimizer.cpp b/csrc/backend_ops/torchscript/optimizer/optimizer.cpp new file mode 100644 index 0000000000..004f9969c8 --- /dev/null +++ b/csrc/backend_ops/torchscript/optimizer/optimizer.cpp @@ -0,0 +1,31 @@ +#include "optimizer.h" + +#include +#include + +#if TORCH_VERSION_MINOR >= 9 +#include +#include +#include +#endif + +namespace mmdeploy { +Module optimize_for_torchscript(const Module& model) { + auto frozen_model = freeze_module(model); + auto graph = frozen_model.get_method("forward").graph(); + OptimizeFrozenGraph(graph, true); + +#if TORCH_VERSION_MINOR >= 9 + FuseFrozenConvAddRelu(graph); + ConvertFrozenOpsToMKLDNN(graph); + FrozenLinearTranspose(graph); +#endif + + // TODO: add more custom passes + + return frozen_model; +} + +// TODO: add optimizer for other backend/onnx + +} // namespace mmdeploy \ No newline at end of file diff --git a/csrc/backend_ops/torchscript/optimizer/optimizer.h b/csrc/backend_ops/torchscript/optimizer/optimizer.h new file mode 100644 index 0000000000..7246b2d83c --- /dev/null +++ b/csrc/backend_ops/torchscript/optimizer/optimizer.h @@ -0,0 +1,7 @@ +#include + +namespace mmdeploy { +using torch::jit::script::Module; + +Module optimize_for_torchscript(const Module &model); +} // namespace mmdeploy From a2b69fc9f6a04c1027526685b1f6bb3dc55c1a59 Mon Sep 17 00:00:00 2001 From: grimoire Date: Fri, 28 Jan 2022 19:43:23 +0800 Subject: [PATCH 2/2] add python api --- .../torchscript/optimizer/CMakeLists.txt | 1 - csrc/backend_ops/torchscript/optimizer/main.cpp | 2 +- .../torchscript/optimizer/optimizer.cpp | 2 +- mmdeploy/apis/pytorch2torchscript.py | 17 +++++++++++++---- mmdeploy/backend/torchscript/__init__.py | 4 ++-- mmdeploy/backend/torchscript/init_plugins.py | 17 +++++++++++++++++ 6 files changed, 34 insertions(+), 9 deletions(-) diff --git a/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt b/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt index 3a0121bc82..cacfece52f 100644 --- a/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt +++ b/csrc/backend_ops/torchscript/optimizer/CMakeLists.txt @@ -9,4 +9,3 @@ file(GLOB_RECURSE OPTIMIZER_SRCS *.cpp) add_executable(${PROJECT_NAME} ${OPTIMIZER_SRCS}) target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES}) target_link_directories(${PROJECT_NAME} PRIVATE mmdeploy::torchscript_ops) - diff --git a/csrc/backend_ops/torchscript/optimizer/main.cpp b/csrc/backend_ops/torchscript/optimizer/main.cpp index cfd69d025c..89d58dcef6 100644 --- a/csrc/backend_ops/torchscript/optimizer/main.cpp +++ b/csrc/backend_ops/torchscript/optimizer/main.cpp @@ -83,4 +83,4 @@ int main(int argc, char* argv[]) { model.save(output_file); return 0; -} \ No newline at end of file +} diff --git a/csrc/backend_ops/torchscript/optimizer/optimizer.cpp b/csrc/backend_ops/torchscript/optimizer/optimizer.cpp index 004f9969c8..a1c5c0107a 100644 --- a/csrc/backend_ops/torchscript/optimizer/optimizer.cpp +++ b/csrc/backend_ops/torchscript/optimizer/optimizer.cpp @@ -28,4 +28,4 @@ Module optimize_for_torchscript(const Module& model) { // TODO: add optimizer for other backend/onnx -} // namespace mmdeploy \ No newline at end of file +} // namespace mmdeploy diff --git a/mmdeploy/apis/pytorch2torchscript.py b/mmdeploy/apis/pytorch2torchscript.py index 92835fd059..2d1a1f2543 100644 --- a/mmdeploy/apis/pytorch2torchscript.py +++ b/mmdeploy/apis/pytorch2torchscript.py @@ -1,12 +1,14 @@ import os.path as osp +from subprocess import call from typing import Any, Optional, Union import mmcv import torch -from mmdeploy.backend.torchscript import get_ops_path +from mmdeploy.backend.torchscript import get_ops_path, get_optimizer_path from mmdeploy.core import RewriterContext, patch_model -from mmdeploy.utils import get_backend, get_input_shape, load_config +from mmdeploy.utils import (get_backend, get_input_shape, get_root_logger, + load_config) def torch2torchscript_impl(model: torch.nn.Module, input: torch.Tensor, @@ -39,10 +41,17 @@ def torch2torchscript_impl(model: torch.nn.Module, input: torch.Tensor, True): ts_model = torch.jit.trace(patched_model, input) - # TODO: custom optimize - + # save model torch.jit.save(ts_model, output_file) + # perform optimize + optimizers_path = get_optimizer_path() + if len(optimizers_path) > 0: + # optimize model + logger = get_root_logger() + logger.info('perform torchscript optimizer.') + call([optimizers_path, '-backend=' + backend, output_file]) + def torch2torchscript(img: Any, work_dir: str, diff --git a/mmdeploy/backend/torchscript/__init__.py b/mmdeploy/backend/torchscript/__init__.py index c335a7f171..6557455d69 100644 --- a/mmdeploy/backend/torchscript/__init__.py +++ b/mmdeploy/backend/torchscript/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. # flake8: noqa -from .init_plugins import get_ops_path +from .init_plugins import get_ops_path, get_optimizer_path def is_available(): @@ -13,7 +13,7 @@ def is_available(): return True -__all__ = ['get_ops_path'] +__all__ = ['get_ops_path', 'get_optimizer_path'] if is_available(): from .wrapper import TorchscriptWrapper diff --git a/mmdeploy/backend/torchscript/init_plugins.py b/mmdeploy/backend/torchscript/init_plugins.py index 3a94e791f3..7c4aff9190 100644 --- a/mmdeploy/backend/torchscript/init_plugins.py +++ b/mmdeploy/backend/torchscript/init_plugins.py @@ -1,5 +1,6 @@ import glob import os.path as osp +import platform def get_ops_path() -> str: @@ -16,3 +17,19 @@ def get_ops_path() -> str: paths = glob.glob(wildcard) lib_path = paths[0] if len(paths) > 0 else '' return lib_path + + +def get_optimizer_path() -> str: + """Get ts_optimizer path. + + Returns: + str: A path of ts_optimizer tool. + """ + wildcard = osp.abspath( + osp.join(osp.dirname(__file__), '../../../build/bin/ts_optimizer')) + if platform.system() == 'Windows': + wildcard += '.exe' + + paths = glob.glob(wildcard) + lib_path = paths[0] if len(paths) > 0 else '' + return lib_path