support coreml (#760)

* sdk inference * fix typo * fix typo * add convert things * fix missling name * add cls support * add more pytorch rewriter * add det support * support det wip * make Model export model_path * fix nms * add output back * add docstring * fix lint * add coreml build action * add zh docs * add coreml backend check * update ci * update * update * update * update * update * fix lint * update configs * add return value when error occured * update docs * update docs * update docs * fix lint * udpate docs * udpate docs * update Co-authored-by: grimoire <[email protected]>
open-mmlab · Sep 5, 2022 · a0fb3be · a0fb3be
1 parent 3d092bf
commit a0fb3be
Show file tree

Hide file tree

Showing 43 changed files with 1,439 additions and 52 deletions.
diff --git a/.github/workflows/backend-coreml.yml b/.github/workflows/backend-coreml.yml
@@ -0,0 +1,71 @@
+name: backend-coreml
+
+on:
+  push:
+    paths:
+      - "csrc/**"
+      - "demo/csrc/**"
+      - "CMakeLists.txt"
+
+  pull_request:
+    paths:
+      - "csrc/**"
+      - "demo/csrc/**"
+      - "CMakeLists.txt"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+env:
+  DEVELOPER_DIR: /Applications/Xcode_13.4.1.app/Contents/Developer
+permissions:
+  contents: read
+
+jobs:
+  build_macos_arm64:
+    runs-on: macos-12
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          submodules: 'recursive'
+      - name: install opencv
+        run: |
+          wget https://github.com/irexyc/mmdeploy-ci-resource/releases/download/opencv/opencv-osx-arm64-4.6.0.tar.gz
+          mkdir $GITHUB_WORKSPACE/opencv-install
+          tar xf opencv-osx-arm64-4.6.0.tar.gz -C $GITHUB_WORKSPACE/opencv-install
+      - name: install libtorch
+        run: |
+          wget https://github.com/irexyc/mmdeploy-ci-resource/releases/download/libtorch/libtorch-osx-arm64-1.8.0.tar.gz
+          mkdir $GITHUB_WORKSPACE/libtorch-install
+          tar xf libtorch-osx-arm64-1.8.0.tar.gz -C $GITHUB_WORKSPACE/libtorch-install
+      - name: build
+        run: |
+          mkdir build && cd build
+          cmake .. -DCMAKE_OSX_ARCHITECTURES="arm64" \
+            -DCMAKE_SYSTEM_PROCESSOR="arm64" \
+            -DMMDEPLOY_BUILD_SDK=ON \
+            -DMMDEPLOY_TARGET_DEVICES="cpu" \
+            -DMMDEPLOY_CODEBASES=all \
+            -DOpenCV_DIR=$GITHUB_WORKSPACE/opencv-install/lib/cmake/opencv4 \
+            -DTorch_DIR=$GITHUB_WORKSPACE/libtorch-install/share/cmake/Torch \
+            -DMMDEPLOY_TARGET_BACKENDS="coreml" \
+            -DMMDEPLOY_BUILD_EXAMPLES=ON \
+            -DMMDEPLOY_SHARED_LIBS=OFF
+          cmake --build . -j 3
+          cmake --build . --target install
+      - name: build-shared
+        run: |
+          mkdir build-shared && cd build-shared
+          cmake .. -DCMAKE_OSX_ARCHITECTURES="arm64" \
+            -DCMAKE_SYSTEM_PROCESSOR="arm64" \
+            -DMMDEPLOY_BUILD_SDK=ON \
+            -DMMDEPLOY_TARGET_DEVICES="cpu" \
+            -DMMDEPLOY_CODEBASES=all \
+            -DOpenCV_DIR=$GITHUB_WORKSPACE/opencv-install/lib/cmake/opencv4 \
+            -DTorch_DIR=$GITHUB_WORKSPACE/libtorch-install/share/cmake/Torch \
+            -DMMDEPLOY_TARGET_BACKENDS="coreml" \
+            -DMMDEPLOY_BUILD_EXAMPLES=ON \
+            -DMMDEPLOY_SHARED_LIBS=ON
+          cmake --build . -j 3
+          cmake --build . --target install
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -77,6 +77,10 @@ if (MSVC)
     add_compile_options($<$<COMPILE_LANGUAGE:CXX>:/wd4251>)
 endif ()
 
+if(APPLE)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fobjc-arc")
+endif()
+
 add_library(MMDeployStaticModules INTERFACE)
 add_library(MMDeployDynamicModules INTERFACE)
 add_library(MMDeployLibs INTERFACE)

diff --git a/README.md b/README.md
@@ -55,9 +55,9 @@ The currently supported codebases and models are as follows, and more will be in
 
 Models can be exported and run in the following backends, and more will be compatible
 
-| ONNX Runtime | TensorRT | ppl.nn | ncnn | OpenVINO | LibTorch | snpe | Ascend | more                                           |
-| ------------ | -------- | ------ | ---- | -------- | -------- | ---- | ------ | ---------------------------------------------- |
-| ✔️           | ✔️       | ✔️     | ✔️   | ✔️       | ✔️       | ✔️   | ✔️     | [benchmark](docs/en/03-benchmark/benchmark.md) |
+| ONNX Runtime | TensorRT | ppl.nn | ncnn | OpenVINO | LibTorch | snpe | Ascend | Core ML | more                                           |
+| ------------ | -------- | ------ | ---- | -------- | -------- | ---- | ------ | ------- | ---------------------------------------------- |
+| ✔️           | ✔️       | ✔️     | ✔️   | ✔️       | ✔️       | ✔️   | ✔️     | ✔️      | [benchmark](docs/en/03-benchmark/benchmark.md) |
 
 ### Efficient and scalable C/C++ SDK Framework
 
@@ -71,6 +71,7 @@ Please read [getting_started](docs/en/get_started.md) for the basic usage of MMD
   - [Build from Docker](docs/en/01-how-to-build/build_from_docker.md)
   - [Build from Script](docs/en/01-how-to-build/build_from_script.md)
   - [Build for Linux](docs/en/01-how-to-build/linux-x86_64.md)
+  - [Build for macOS](docs/en/01-how-to-build/macos-arm64.md)
   - [Build for Win10](docs/en/01-how-to-build/windows.md)
   - [Build for Android](docs/en/01-how-to-build/android.md)
   - [Build for Jetson](docs/en/01-how-to-build/jetsons.md)

diff --git a/README_zh-CN.md b/README_zh-CN.md
@@ -53,9 +53,9 @@ MMDeploy 是 [OpenMMLab](https://openmmlab.com/) 模型部署工具箱，**为
 
 ### 支持多种推理后端
 
-| ONNX Runtime | TensorRT | ppl.nn | ncnn | OpenVINO | LibTorch | snpe | Ascend | more                                              |
-| ------------ | -------- | ------ | ---- | -------- | -------- | ---- | ------ | ------------------------------------------------- |
-| ✔️           | ✔️       | ✔️     | ✔️   | ✔️       | ✔️       | ✔️   | ✔️     | [benchmark](docs/zh_cn/03-benchmark/benchmark.md) |
+| ONNX Runtime | TensorRT | ppl.nn | ncnn | OpenVINO | LibTorch | snpe | Ascend | Core ML | more                                           |
+| ------------ | -------- | ------ | ---- | -------- | -------- | ---- | ------ | ------- | ---------------------------------------------- |
+| ✔️           | ✔️       | ✔️     | ✔️   | ✔️       | ✔️       | ✔️   | ✔️     | ✔️      | [benchmark](docs/en/03-benchmark/benchmark.md) |
 
 ### SDK 可高度定制化
 
@@ -70,6 +70,7 @@ MMDeploy 是 [OpenMMLab](https://openmmlab.com/) 模型部署工具箱，**为
   - [一键式脚本安装](docs/zh_cn/01-how-to-build/build_from_script.md)
   - [Build from Docker](docs/zh_cn/01-how-to-build/build_from_docker.md)
   - [Build for Linux](docs/zh_cn/01-how-to-build/linux-x86_64.md)
+  - [Build for macOS](docs/zh_cn/01-how-to-build/macos-arm64.md)
   - [Build for Win10](docs/zh_cn/01-how-to-build/windows.md)
   - [Build for Android](docs/zh_cn/01-how-to-build/android.md)
   - [Build for Jetson](docs/en/01-how-to-build/jetsons.md)

diff --git a/configs/_base_/backends/coreml.py b/configs/_base_/backends/coreml.py
@@ -0,0 +1 @@
+backend_config = dict(type='coreml', convert_to='mlprogram')
diff --git a/configs/mmcls/classification_coreml_dynamic-224x224-224x224.py b/configs/mmcls/classification_coreml_dynamic-224x224-224x224.py
@@ -0,0 +1,12 @@
+_base_ = ['../_base_/torchscript_config.py', '../_base_/backends/coreml.py']
+
+codebase_config = dict(type='mmcls', task='Classification')
+
+backend_config = dict(model_inputs=[
+    dict(
+        input_shapes=dict(
+            input=dict(
+                min_shape=[1, 3, 224, 224],
+                max_shape=[8, 3, 224, 224],
+                default_shape=[1, 3, 224, 224])))
+])
diff --git a/configs/mmdet/_base_/base_coreml_static-800x1344.py b/configs/mmdet/_base_/base_coreml_static-800x1344.py
@@ -0,0 +1,11 @@
+_base_ = ['./base_torchscript.py', '../../_base_/backends/coreml.py']
+
+ir_config = dict(input_shape=(1344, 800))
+backend_config = dict(model_inputs=[
+    dict(
+        input_shapes=dict(
+            input=dict(
+                min_shape=[1, 3, 800, 1344],
+                max_shape=[1, 3, 800, 1344],
+                default_shape=[1, 3, 800, 1344])))
+])
diff --git a/configs/mmdet/detection/detection_coreml_static-800x1344.py b/configs/mmdet/detection/detection_coreml_static-800x1344.py
@@ -0,0 +1 @@
+_base_ = ['../_base_/base_coreml_static-800x1344.py']
diff --git a/configs/mmseg/segmentation_coreml_static-512x1024.py b/configs/mmseg/segmentation_coreml_static-512x1024.py
@@ -0,0 +1,14 @@
+_base_ = [
+    '../_base_/torchscript_config.py', '../_base_/backends/coreml.py',
+    './segmentation_static.py'
+]
+
+ir_config = dict(input_shape=[1024, 512])
+backend_config = dict(model_inputs=[
+    dict(
+        input_shapes=dict(
+            input=dict(
+                min_shape=[1, 3, 512, 1024],
+                max_shape=[1, 3, 512, 1024],
+                default_shape=[1, 3, 512, 1024])))
+])
diff --git a/csrc/mmdeploy/backend_ops/CMakeLists.txt b/csrc/mmdeploy/backend_ops/CMakeLists.txt
@@ -32,7 +32,8 @@ if ("ncnn" IN_LIST MMDEPLOY_TARGET_BACKENDS)
 endif ()
 
 # build TorchScript ops
-if ("torchscript" IN_LIST MMDEPLOY_TARGET_BACKENDS)
-  message(STATUS "Build torchsciprt custom ops")
+if ("torchscript" IN_LIST MMDEPLOY_TARGET_BACKENDS
+    OR "coreml" IN_LIST MMDEPLOY_TARGET_BACKENDS)
+  message(STATUS "Build torchscript custom ops")
   add_subdirectory(torchscript)
 endif ()
diff --git a/csrc/mmdeploy/backend_ops/torchscript/bind.cpp b/csrc/mmdeploy/backend_ops/torchscript/bind.cpp
diff --git a/csrc/mmdeploy/backend_ops/torchscript/ops/bind.cpp b/csrc/mmdeploy/backend_ops/torchscript/ops/bind.cpp
@@ -0,0 +1,13 @@
+// Copyright (c) OpenMMLab. All rights reserved.
+#include "torch/script.h"
+
+TORCH_LIBRARY(mmdeploy, m) {
+  m.def(
+       "modulated_deform_conv(Tensor input, Tensor weight, Tensor bias, Tensor offset, Tensor "
+       "mask, "
+       "int kernel_h, int kernel_w, int stride_h, int stride_w, int pad_h, int pad_w, int "
+       "dilation_h,int dilation_w, int groups, int deform_groups, bool with_bias) -> Tensor")
+      .def(
+          "coreml_nms(Tensor boxes, Tensor scores, float iou_threshold, "
+          "float score_threshold, int max_boxes) -> Tensor[]");
+}
diff --git a/csrc/mmdeploy/backend_ops/torchscript/ops/coreml_nms/coreml_nms_cpu.cpp b/csrc/mmdeploy/backend_ops/torchscript/ops/coreml_nms/coreml_nms_cpu.cpp
@@ -0,0 +1,31 @@
+
+#include <assert.h>
+
+#include <vector>
+
+#include "torch/script.h"
+namespace mmdeploy {
+
+using at::Tensor;
+
+std::vector<Tensor> coreml_nms_cpu(Tensor boxes, Tensor scores, double iou_threshold,
+                                   double score_threshold, int64_t max_boxes) {
+  assert(boxes.dim() == 3);  // bboxes with shape (batch_size, num_bboxes, 4)
+  assert(boxes.size(2) == 4);
+  assert(boxes.size(0) == scores.size(0));  // check batch size
+  assert(boxes.size(1) == scores.size(1));  // check num boxes
+
+  auto batch_size = boxes.size(0);
+  auto num_boxes = boxes.size(1);
+  auto num_classes = scores.size(2);
+
+  Tensor ret_boxes = at::zeros({batch_size, max_boxes, 4});
+  Tensor ret_scores = at::zeros({batch_size, max_boxes, num_classes});
+  Tensor indices = at::zeros({batch_size, max_boxes}, at::kInt);
+  Tensor num_outputs = at::zeros({batch_size}, at::kInt);
+
+  return std::vector<Tensor>({ret_boxes, ret_scores, indices, num_outputs});
+}
+
+TORCH_LIBRARY_IMPL(mmdeploy, CPU, m) { m.impl("coreml_nms", coreml_nms_cpu); }
+}  // namespace mmdeploy
diff --git a/csrc/mmdeploy/core/model.cpp b/csrc/mmdeploy/core/model.cpp
@@ -20,6 +20,7 @@ Model::Model(const std::string& model_path) {
 Model::Model(const void* buffer, size_t size) { Init(buffer, size).value(); }
 
 Result<void> Model::Init(const std::string& model_path) {
+  model_path_ = model_path;
   if (!fs::exists(model_path)) {
     MMDEPLOY_ERROR("'{}' doesn't exist", model_path);
     return Status(eFileNotExist);
@@ -45,6 +46,8 @@ Result<void> Model::Init(const std::string& model_path) {
   return Status(eNotSupported);
 }
 
+const std::string& Model::GetModelPath() const { return model_path_; }
+
 Result<void> Model::Init(const void* buffer, size_t size) {
   auto registry = ModelRegistry::Get();
   auto entries = registry.ListEntries();

diff --git a/csrc/mmdeploy/core/model.h b/csrc/mmdeploy/core/model.h
@@ -94,7 +94,14 @@ class MMDEPLOY_API Model {
    */
   explicit operator bool() const { return impl_ != nullptr; }
 
+  /**
+   * @brief get model_path that init with DirectoryModel
+   * @return file path of an sdk model
+   */
+  const std::string& GetModelPath() const;
+
  private:
+  std::string model_path_;
   std::shared_ptr<ModelImpl> impl_;
   deploy_meta_info_t meta_;
 };

diff --git a/csrc/mmdeploy/net/CMakeLists.txt b/csrc/mmdeploy/net/CMakeLists.txt
@@ -34,5 +34,9 @@ if ("torchscript" IN_LIST MMDEPLOY_TARGET_BACKENDS)
     add_subdirectory(torchscript)
 endif ()
 
+if ("coreml" IN_LIST MMDEPLOY_TARGET_BACKENDS)
+    add_subdirectory(coreml)
+endif ()
+
 mmdeploy_add_module(${PROJECT_NAME} net_module.cpp)
 add_library(mmdeploy::net_module ALIAS ${PROJECT_NAME})
diff --git a/csrc/mmdeploy/net/coreml/CMakeLists.txt b/csrc/mmdeploy/net/coreml/CMakeLists.txt
@@ -0,0 +1,14 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+project(mmdeploy_coreml_net)
+
+if ("cpu" IN_LIST MMDEPLOY_TARGET_DEVICES)
+    find_library(CORE_ML CoreML)
+    find_library(FOUNDATION Foundation)
+    mmdeploy_add_module(${PROJECT_NAME} coreml_net.mm)
+    target_include_directories(${PROJECT_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+    target_link_libraries(${PROJECT_NAME} PRIVATE ${CORE_ML} ${FOUNDATION})
+    add_library(mmdeploy::coreml_net ALIAS ${PROJECT_NAME})
+else ()
+    message(ERROR "'coreml_net' is NOT supported in target devices: ${MMDEPLOY_TARGET_DEVICES}")
+endif ()
diff --git a/csrc/mmdeploy/net/coreml/coreml_net.h b/csrc/mmdeploy/net/coreml/coreml_net.h
@@ -0,0 +1,37 @@
+// Copyright (c) OpenMMLab. All rights reserved.
+
+#ifndef MMDEPLOY_SRC_NET_COREML_COREML_NET_H_
+#define MMDEPLOY_SRC_NET_COREML_COREML_NET_H_
+
+#include "mmdeploy/core/net.h"
+
+namespace mmdeploy {
+
+namespace coreml {
+class Execution;
+}  // namespace coreml
+
+class CoreMLNet : public Net {
+ public:
+  ~CoreMLNet() override = default;
+  Result<void> Init(const Value& cfg) override;
+  Result<void> Deinit() override;
+  Result<Span<Tensor>> GetInputTensors() override;
+  Result<Span<Tensor>> GetOutputTensors() override;
+  Result<void> Reshape(Span<TensorShape> input_shapes) override;
+  Result<void> Forward() override;
+  Result<void> ForwardAsync(Event* event) override;
+
+ private:
+  std::unique_ptr<coreml::Execution> execution_;
+  std::vector<Tensor> input_tensors_;
+  std::vector<Tensor> output_tensors_;
+  Device device_;
+  Stream stream_;
+
+  friend class coreml::Execution;
+};
+
+}  // namespace mmdeploy
+
+#endif  // MMDEPLOY_SRC_NET_ORT_ORT_NET_H_
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		backend_config = dict(type='coreml', convert_to='mlprogram')
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		_base_ = ['../_base_/base_coreml_static-800x1344.py']