From 3a6def6fee479ef42bc878a647bd6bcf390b9b3f Mon Sep 17 00:00:00 2001 From: liuyulvv Date: Fri, 5 Aug 2022 18:24:12 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0Windows=20onnxruntime=20cuda?= =?UTF-8?q?=E6=8E=A8=E7=90=86=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- WORKSPACE | 6 + mediapipe/calculators/tensor/BUILD | 32 ++++ .../calculators/tensor/inference_calculator.h | 4 + .../tensor/inference_calculator_onnx_cuda.cc | 150 ++++++++++++++++++ third_party/onnxruntime_windows.BUILD | 17 ++ 5 files changed, 209 insertions(+) create mode 100644 mediapipe/calculators/tensor/inference_calculator_onnx_cuda.cc create mode 100644 third_party/onnxruntime_windows.BUILD diff --git a/WORKSPACE b/WORKSPACE index 336783b67..bb80da9a1 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -226,6 +226,12 @@ new_local_repository( path = "D:\\opencv\\build", ) +new_local_repository( + name = "windows_onnxruntime", + build_file = "@//third_party:onnxruntime_windows.BUILD", + path = "D:\\onnxruntime\\onnxruntime-win-x64-gpu-1.12.0", +) + http_archive( name = "android_opencv", build_file = "@//third_party:opencv_android.BUILD", diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD index c95938ab6..55651107a 100644 --- a/mediapipe/calculators/tensor/BUILD +++ b/mediapipe/calculators/tensor/BUILD @@ -245,6 +245,38 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "inference_calculator_onnx_cuda", + srcs = [ + "inference_calculator_onnx_cuda.cc", + ], + copts = select({ + # TODO: fix tensor.h not to require this, if possible + "//mediapipe:apple": [ + "-x objective-c++", + "-fobjc-arc", # enable reference-counting + ], + "//conditions:default": [], + }), + visibility = ["//visibility:public"], + deps = [ + ":inference_calculator_interface", + "@com_google_absl//absl/memory", + "@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate", + "@org_tensorflow//tensorflow/lite:framework_stable", + "@org_tensorflow//tensorflow/lite/c:c_api_types", + "@windows_onnxruntime//:onnxruntime", + ] + select({ + "//conditions:default": [ + "//mediapipe/util:cpu_util", + ], + }) + select({ + "//conditions:default": [], + "//mediapipe:android": ["@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate"], + }), + alwayslink = 1, +) + cc_library( name = "inference_calculator_gl_if_compute_shader_available", visibility = ["//visibility:public"], diff --git a/mediapipe/calculators/tensor/inference_calculator.h b/mediapipe/calculators/tensor/inference_calculator.h index db5fa4f5f..126567de0 100644 --- a/mediapipe/calculators/tensor/inference_calculator.h +++ b/mediapipe/calculators/tensor/inference_calculator.h @@ -145,6 +145,10 @@ struct InferenceCalculatorCpu : public InferenceCalculator { static constexpr char kCalculatorName[] = "InferenceCalculatorCpu"; }; +struct InferenceCalculatorOnnxCUDA : public InferenceCalculator { + static constexpr char kCalculatorName[] = "InferenceCalculatorOnnxCUDA"; +}; + } // namespace api2 } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_calculator_onnx_cuda.cc b/mediapipe/calculators/tensor/inference_calculator_onnx_cuda.cc new file mode 100644 index 000000000..d9bcf384a --- /dev/null +++ b/mediapipe/calculators/tensor/inference_calculator_onnx_cuda.cc @@ -0,0 +1,150 @@ +// Copyright 2019 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "absl/memory/memory.h" +#include "mediapipe/calculators/tensor/inference_calculator.h" +#include "onnxruntime_cxx_api.h" +#include "tensorflow/lite/c/c_api_types.h" +#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h" +#include "tensorflow/lite/interpreter_builder.h" +#include +#include +#include +#include + +namespace mediapipe { +namespace api2 { + +namespace { + +int64_t value_size_of(const std::vector& dims) { + if (dims.empty()) return 0; + int64_t value_size = 1; + for (const auto& size : dims) value_size *= size; + return value_size; +} + +} // namespace + +class InferenceCalculatorOnnxCUDAImpl + : public NodeImpl { +public: + static absl::Status UpdateContract(CalculatorContract* cc); + + absl::Status Open(CalculatorContext* cc) override; + absl::Status Process(CalculatorContext* cc) override; + absl::Status Close(CalculatorContext* cc) override; + +private: + absl::Status LoadModel(const std::string& path); + + Ort::Env env_; + std::unique_ptr session_; + Ort::AllocatorWithDefaultOptions allocator; + Ort::MemoryInfo memory_info_handler = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + std::vector m_input_names; + std::vector m_output_names; +}; + +absl::Status InferenceCalculatorOnnxCUDAImpl::UpdateContract( + CalculatorContract* cc) { + const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>(); + RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected()) + << "Either model as side packet or model path in options is required."; + return absl::OkStatus(); +} + +absl::Status InferenceCalculatorCpuImpl::LoadModel(const std::string& path) { + auto model_path = std::wstring(path.begin(), path.end()); + Ort::SessionOptions session_options; + OrtCUDAProviderOptions cuda_options; + session_options.AppendExecutionProvider_CUDA(cuda_options); + session_ = std::make_unique(env_, model_path.c_str(), session_options); + size_t num_input_nodes = session_->GetInputCount(); + size_t num_output_nodes = session_->GetOutputCount(); + m_input_names.reserve(num_input_nodes); + m_output_names.reserve(num_output_nodes); + for (int i = 0; i < num_input_nodes; i++) { + char* input_name = session_->GetInputName(i, allocator); + m_input_names.push_back(input_name); + } + for (int i = 0; i < num_output_nodes; i++) { + char* output_name = session_->GetOutputName(i, allocator); + m_output_names.push_back(output_name); + } + return absl::OkStatus(); +} + +absl::Status InferenceCalculatorOnnxCUDAImpl::Open(CalculatorContext* cc) { + const auto& options = cc->Options(); + if (!options.model_path().empty()) { + return LoadModel(options.model_path()); + } + if (!options.landmark_path().empty()) { + return LoadModel(options.landmark_path()); + } + return absl::Status(mediapipe::StatusCode::kNotFound, "Must specify Onnx model path."); +} + +absl::Status InferenceCalculatorOnnxCUDAImpl::Process(CalculatorContext* cc) { + if (kInTensors(cc).IsEmpty()) { + return absl::OkStatus(); + } + const auto& input_tensors = *kInTensors(cc); + RET_CHECK(!input_tensors.empty()); + auto input_tensor_type = int(input_tensors[0].element_type()); + std::vector ort_input_tensors; + ort_input_tensors.reserve(input_tensors.size()); + for (const auto& tensor : input_tensors) { + auto& inputDims = tensor.shape().dims; + std::vector src_dims{inputDims[0], inputDims[1], inputDims[2], inputDims[3]}; + auto src_value_size = value_size_of(src_dims); + auto input_tensor_view = tensor.GetCpuReadView(); + auto input_tensor_buffer = const_cast(input_tensor_view.buffer()); + auto tmp_tensor = Ort::Value::CreateTensor(memory_info_handler, input_tensor_buffer, src_value_size, src_dims.data(), src_dims.size()); + ort_input_tensors.emplace_back(std::move(tmp_tensor)); + } + auto output_tensors = absl::make_unique>(); + std::vector onnx_output_tensors; + try { + onnx_output_tensors = session_->Run( + Ort::RunOptions{nullptr}, m_input_names.data(), + ort_input_tensors.data(), ort_input_tensors.size(), m_output_names.data(), + m_output_names.size()); + } catch (Ort::Exception& e) { + LOG(ERROR) << "Run error msg:" << e.what(); + } + for (const auto& tensor : onnx_output_tensors) { + auto info = tensor.GetTensorTypeAndShapeInfo(); + auto dims = info.GetShape(); + std::vector tmp_dims; + for (const auto& i : dims) { + tmp_dims.push_back(i); + } + output_tensors->emplace_back(Tensor::ElementType::kFloat32, Tensor::Shape{tmp_dims}); + auto cpu_view = output_tensors->back().GetCpuWriteView(); + std::memcpy(cpu_view.buffer(), tensor.GetTensorData(), output_tensors->back().bytes()); + } + kOutTensors(cc).Send(std::move(output_tensors)); + return absl::OkStatus(); +} + +absl::Status InferenceCalculatorOnnxCUDAImpl::Close(CalculatorContext* cc) { + interpreter_ = nullptr; + delegate_ = nullptr; + return absl::OkStatus(); +} + +} // namespace api2 +} // namespace mediapipe diff --git a/third_party/onnxruntime_windows.BUILD b/third_party/onnxruntime_windows.BUILD new file mode 100644 index 000000000..9ba7e286e --- /dev/null +++ b/third_party/onnxruntime_windows.BUILD @@ -0,0 +1,17 @@ +cc_library( + name = "onnxruntime", + srcs = [ + "lib/onnxruntime.dll", + "lib/onnxruntime.lib", + "lib/onnxruntime_providers_cuda.dll", + "lib/onnxruntime_providers_cuda.lib", + "lib/onnxruntime_providers_shared.dll", + "lib/onnxruntime_providers_shared.lib", + "lib/onnxruntime_providers_tensorrt.dll", + "lib/onnxruntime_providers_tensorrt.lib", + ], + hdrs = glob(["include/*.h*"]), + includes = ["include/"], + linkstatic = 1, + visibility = ["//visibility:public"], +)