From f3ddc40ca44e35d42b39cdb9a72387d9d22127aa Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Tue, 26 Mar 2019 20:57:18 -0700 Subject: [PATCH] Move weight offload inside backend construction functor (#18385) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/18385 By moving the weight offload into the backend initialization function, we can instantiate the backend once by creating the OnnxifiOp once and then clean up the parameter workspace. And we need to keep hold of that instantiated net (OnnxifiOp) without cleaning it. Subsequent ctor of OnnxifiOp of the same model will hit the cached backend and they will not look into weight offloading, which is safe as the weight is already gone. Reviewed By: ipiszy Differential Revision: D14590379 fbshipit-source-id: f7f34016e09777ad3df0af487885cd14658e1044 --- caffe2/operators/onnxifi_op.cc | 13 +++++---- caffe2/operators/onnxifi_op.h | 64 +++++++++++++++++++++--------------------- 2 files changed, 39 insertions(+), 38 deletions(-) diff --git a/caffe2/operators/onnxifi_op.cc b/caffe2/operators/onnxifi_op.cc index 88e3662..e8b582f 100644 --- a/caffe2/operators/onnxifi_op.cc +++ b/caffe2/operators/onnxifi_op.cc @@ -91,9 +91,11 @@ template <> std::vector OnnxifiOp::buildInitializationList( Workspace* ws, - std::unordered_set* initialization_list, + const std::vector& initializers, std::vector* weight_names, std::vector>* weight_shapes) { + std::unordered_set initialization_list( + initializers.begin(), initializers.end()); const std::vector& ws_blobs = ws->Blobs(); // Since onnxTensorDescriptorV1.name will point into the memory in // weight_names, we need to prevent weight_names from reallocating by @@ -101,18 +103,17 @@ OnnxifiOp::buildInitializationList( weight_names->reserve(ws_blobs.size()); std::vector descs; for (const auto& s : ws_blobs) { - auto it = initialization_list->find(s); - if (it != initialization_list->end()) { + auto it = initialization_list.find(s); + if (it != initialization_list.end()) { weight_names->emplace_back(s); onnxTensorDescriptorV1 tensor_desc; tensor_desc.name = weight_names->back().c_str(); BlobToTensorDescriptor(s, ws, &tensor_desc, weight_shapes); descs.push_back(tensor_desc); - initialization_list->erase(it); + initialization_list.erase(it); } } - CAFFE_ENFORCE( - initialization_list->empty(), "Unfulfilled initialization list"); + CAFFE_ENFORCE(initialization_list.empty(), "Unfulfilled initialization list"); return descs; } diff --git a/caffe2/operators/onnxifi_op.h b/caffe2/operators/onnxifi_op.h index eb9c44a..3965aaf 100644 --- a/caffe2/operators/onnxifi_op.h +++ b/caffe2/operators/onnxifi_op.h @@ -68,24 +68,15 @@ class OnnxifiOp final : public Operator { std::vector property_pointers; std::vector int_args; std::vector float_args; - BuildPropertyList(operator_def, &property_pointers, &int_args, &float_args); - - // Pull the weights from workspace and feed it to the backend through - // setGraphIO. Notice that since we may have rewritten the net, we need to - // map the weight names - auto initializers = - this->template GetRepeatedArgument("initializers"); - std::unordered_set initializer_set; - for (auto it = initializers.begin(); it != initializers.end(); ++it) { - auto key = *it; - initializer_set.emplace(key); - } - std::vector weight_names; - std::vector> weight_shapes; - auto weight_descs = buildInitializationList( - ws, &initializer_set, &weight_names, &weight_shapes); - - BuildBackendAndGraph(property_pointers, onnx_model_str, weight_descs); + buildPropertyList(operator_def, &property_pointers, &int_args, &float_args); + + // Initialize the backend if it has not been already created. When we + // initialized the backend, we will get the weights (initializers) from the + // workspace and offload onto the backend. This should be done only once. + // Subsequent call of this function with the same model id should find a + // cached backend and therefore there is no need to repeat the above + // process. + buildBackendAndGraph(ws, property_pointers, onnx_model_str); } ~OnnxifiOp() { @@ -109,7 +100,7 @@ class OnnxifiOp final : public Operator { return type; } - void BuildPropertyList( + void buildPropertyList( const OperatorDef& /* unused */, std::vector* property_list, std::vector* /* unused */, @@ -117,26 +108,28 @@ class OnnxifiOp final : public Operator { property_list->push_back(ONNXIFI_BACKEND_PROPERTY_NONE); } - void BuildBackendAndGraph( + void buildBackendAndGraph( + Workspace* ws, const std::vector& property_pointers, - const std::string& onnx_model_str, - const std::vector& weight_descs) { + const std::string& onnx_model_str) { op_id_string_ = this->template GetSingleArgument("model_id", "") + ":" + this->template GetSingleArgument("net_pos", ""); + auto initializers = + this->template GetRepeatedArgument("initializers"); // Build the Onnxifi engine auto backend_index = this->template GetSingleArgument("backend_id", 0); - onnxifi_library* lib = lib_; - auto creator = [lib, + auto creator = [this, + ws, property_pointers, backend_index, &onnx_model_str, - &weight_descs]() { + &initializers]() { std::vector backend_ids; size_t num_backends{0}; CAFFE_ENFORCE_EQ( - lib->onnxGetBackendIDs(nullptr, &num_backends), + lib_->onnxGetBackendIDs(nullptr, &num_backends), ONNXIFI_STATUS_FALLBACK); CAFFE_ENFORCE_GT( num_backends, 0, "At least 1 onnxifi backend should be available"); @@ -149,14 +142,14 @@ class OnnxifiOp final : public Operator { num_backends); backend_ids.resize(num_backends); CAFFE_ENFORCE_EQ( - lib->onnxGetBackendIDs(backend_ids.data(), &num_backends), + lib_->onnxGetBackendIDs(backend_ids.data(), &num_backends), ONNXIFI_STATUS_SUCCESS); onnxBackendID backend_id = backend_ids[backend_index]; onnxBackend backend{nullptr}; CAFFE_ENFORCE_EQ( - lib->onnxInitBackend(backend_id, property_pointers.data(), &backend), + lib_->onnxInitBackend(backend_id, property_pointers.data(), &backend), ONNXIFI_STATUS_SUCCESS); // Release unused backend ids. @@ -164,11 +157,18 @@ class OnnxifiOp final : public Operator { if (i == backend_index) { continue; } - lib->onnxReleaseBackendID(backend_ids[i]); + lib_->onnxReleaseBackendID(backend_ids[i]); } + + // Get weights + std::vector weight_names; + std::vector> weight_shapes; + auto weight_descs = buildInitializationList( + ws, initializers, &weight_names, &weight_shapes); + onnxGraph graph{nullptr}; CAFFE_ENFORCE_EQ( - lib->onnxInitGraph( + lib_->onnxInitGraph( backend, nullptr, onnx_model_str.size(), @@ -179,7 +179,7 @@ class OnnxifiOp final : public Operator { ONNXIFI_STATUS_SUCCESS); return std::make_shared( - backend_id, backend, graph, lib); + backend_id, backend, graph, lib_); }; backend_graph_shared_ptr_ = backend_graph_map_ptr_->insert(op_id_string_, creator); @@ -204,7 +204,7 @@ class OnnxifiOp final : public Operator { std::vector buildInitializationList( Workspace* ws, - std::unordered_set* initialization_list, + const std::vector& initializers, std::vector* weight_names, std::vector>* weight_shapes); -- 2.7.4