diff --git a/apps/android_camera/models/prepare_model.py b/apps/android_camera/models/prepare_model.py index ab20e028c2ad..f155d46c31a4 100644 --- a/apps/android_camera/models/prepare_model.py +++ b/apps/android_camera/models/prepare_model.py @@ -106,7 +106,7 @@ def main(model_str, output_path): f.write(graph) print("dumping params...") with open(output_path_str + "/" + "deploy_param.params", "wb") as f: - f.write(relay.save_param_dict(params)) + f.write(runtime.save_param_dict(params)) print("dumping labels...") synset_url = "".join( [ diff --git a/apps/bundle_deploy/build_model.py b/apps/bundle_deploy/build_model.py index 0991ac9ad94b..8fbc01bcf4a6 100644 --- a/apps/bundle_deploy/build_model.py +++ b/apps/bundle_deploy/build_model.py @@ -20,7 +20,7 @@ import os from tvm import relay import tvm -from tvm import te +from tvm import te, runtime import logging import json from tvm.contrib import cc as _cc @@ -70,7 +70,7 @@ def build_module(opts): with open( os.path.join(build_dir, file_format_str.format(name="params", ext="bin")), "wb" ) as f_params: - f_params.write(relay.save_param_dict(params)) + f_params.write(runtime.save_param_dict(params)) def build_test_module(opts): @@ -113,7 +113,7 @@ def build_test_module(opts): with open( os.path.join(build_dir, file_format_str.format(name="test_params", ext="bin")), "wb" ) as f_params: - f_params.write(relay.save_param_dict(lowered_params)) + f_params.write(runtime.save_param_dict(lowered_params)) with open( os.path.join(build_dir, file_format_str.format(name="test_data", ext="bin")), "wb" ) as fp: diff --git a/apps/bundle_deploy/runtime.cc b/apps/bundle_deploy/runtime.cc index 3224028b60a1..2f7e3848b4bf 100644 --- a/apps/bundle_deploy/runtime.cc +++ b/apps/bundle_deploy/runtime.cc @@ -23,6 +23,7 @@ #include #include "../../src/runtime/c_runtime_api.cc" +#include "../../src/runtime/container.cc" #include "../../src/runtime/cpu_device_api.cc" #include "../../src/runtime/file_utils.cc" #include "../../src/runtime/graph/graph_runtime.cc" diff --git a/apps/sgx/src/build_model.py b/apps/sgx/src/build_model.py index 868d3bcb9fc4..1fc297d8a094 100755 --- a/apps/sgx/src/build_model.py +++ b/apps/sgx/src/build_model.py @@ -23,7 +23,7 @@ from os import path as osp import sys -from tvm import relay +from tvm import relay, runtime from tvm.relay import testing import tvm from tvm import te @@ -49,7 +49,7 @@ def main(): with open(osp.join(build_dir, "graph.json"), "w") as f_graph_json: f_graph_json.write(graph) with open(osp.join(build_dir, "params.bin"), "wb") as f_params: - f_params.write(relay.save_param_dict(params)) + f_params.write(runtime.save_param_dict(params)) if __name__ == "__main__": diff --git a/apps/wasm-standalone/wasm-graph/tools/build_graph_lib.py b/apps/wasm-standalone/wasm-graph/tools/build_graph_lib.py index 42695d28fadb..3d8a349b8744 100644 --- a/apps/wasm-standalone/wasm-graph/tools/build_graph_lib.py +++ b/apps/wasm-standalone/wasm-graph/tools/build_graph_lib.py @@ -24,7 +24,7 @@ import onnx import tvm -from tvm import relay +from tvm import relay, runtime def _get_mod_and_params(model_file): @@ -60,7 +60,7 @@ def build_graph_lib(model_file, opt_level): f_graph.write(graph_json) with open(os.path.join(out_dir, "graph.params"), "wb") as f_params: - f_params.write(relay.save_param_dict(params)) + f_params.write(runtime.save_param_dict(params)) if __name__ == "__main__": diff --git a/docs/deploy/android.rst b/docs/deploy/android.rst index 8c8fcfb49679..256978d00607 100644 --- a/docs/deploy/android.rst +++ b/docs/deploy/android.rst @@ -31,7 +31,7 @@ The code below will save the compilation output which is required on android tar with open("deploy_graph.json", "w") as fo: fo.write(graph.json()) with open("deploy_param.params", "wb") as fo: - fo.write(relay.save_param_dict(params)) + fo.write(runtime.save_param_dict(params)) deploy_lib.so, deploy_graph.json, deploy_param.params will go to android target. diff --git a/golang/sample/gen_mobilenet_lib.py b/golang/sample/gen_mobilenet_lib.py index b82e0c476b9f..12f215b4fd9c 100644 --- a/golang/sample/gen_mobilenet_lib.py +++ b/golang/sample/gen_mobilenet_lib.py @@ -16,7 +16,7 @@ # under the License. import os -from tvm import relay, transform +from tvm import relay, transform, runtime from tvm.contrib.download import download_testdata @@ -94,4 +94,4 @@ def extract(path): fo.write(graph) with open("./mobilenet.params", "wb") as fo: - fo.write(relay.save_param_dict(params)) + fo.write(runtime.save_param_dict(params)) diff --git a/python/tvm/contrib/debugger/debug_result.py b/python/tvm/contrib/debugger/debug_result.py index 3159ab34397a..f58947f0766f 100644 --- a/python/tvm/contrib/debugger/debug_result.py +++ b/python/tvm/contrib/debugger/debug_result.py @@ -264,8 +264,4 @@ def save_tensors(params): """ _save_tensors = tvm.get_global_func("tvm.relay._save_param_dict") - args = [] - for k, v in params.items(): - args.append(k) - args.append(tvm.nd.array(v)) - return _save_tensors(*args) + return _save_tensors(params) diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index fc1805ee0ab4..83791e50f6d5 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -24,7 +24,7 @@ import tvm from tvm import autotvm, auto_scheduler -from tvm import relay +from tvm import relay, runtime from tvm.contrib import cc from tvm.contrib import utils @@ -282,7 +282,7 @@ def save_module(module_path, graph, lib, params, cross=None): with open(temp.relpath(param_name), "wb") as params_file: logger.debug("writing params to file to %s", params_file.name) - params_file.write(relay.save_param_dict(params)) + params_file.write(runtime.save_param_dict(params)) logger.debug("saving module as tar file to %s", module_path) with tarfile.open(module_path, "w") as tar: diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py index 87ea3be1436a..1d23ccfb0c00 100644 --- a/python/tvm/driver/tvmc/runner.py +++ b/python/tvm/driver/tvmc/runner.py @@ -24,11 +24,11 @@ import tempfile import numpy as np -import tvm from tvm import rpc from tvm.autotvm.measure import request_remote from tvm.contrib import graph_runtime as runtime from tvm.contrib.debugger import debug_runtime +from tvm.relay import load_param_dict from . import common from .common import TVMCException @@ -163,9 +163,8 @@ def get_input_info(graph_str, params): shape_dict = {} dtype_dict = {} - # Use a special function to load the binary params back into a dict - load_arr = tvm.get_global_func("tvm.relay._load_param_dict")(params) - param_names = [v.name for v in load_arr] + params_dict = load_param_dict(params) + param_names = [k for (k, v) in params_dict.items()] graph = json.loads(graph_str) for node_id in graph["arg_nodes"]: node = graph["nodes"][node_id] diff --git a/python/tvm/relay/param_dict.py b/python/tvm/relay/param_dict.py index 2d0398e20486..2714607947f3 100644 --- a/python/tvm/relay/param_dict.py +++ b/python/tvm/relay/param_dict.py @@ -16,12 +16,7 @@ # under the License. # pylint: disable=invalid-name """Helper utility to save parameter dicts.""" -import tvm -import tvm._ffi - - -_save_param_dict = tvm._ffi.get_global_func("tvm.relay._save_param_dict") -_load_param_dict = tvm._ffi.get_global_func("tvm.relay._load_param_dict") +import tvm.runtime def save_param_dict(params): @@ -30,6 +25,9 @@ def save_param_dict(params): The result binary bytes can be loaded by the GraphModule with API "load_params". + .. deprecated:: 0.9.0 + Use :py:func:`tvm.runtime.save_param_dict` instead. + Parameters ---------- params : dict of str to NDArray @@ -47,21 +45,20 @@ def save_param_dict(params): # set up the parameter dict params = {"param0": arr0, "param1": arr1} # save the parameters as byte array - param_bytes = tvm.relay.save_param_dict(params) + param_bytes = tvm.runtime.save_param_dict(params) # We can serialize the param_bytes and load it back later. # Pass in byte array to module to directly set parameters - graph_runtime_mod.load_params(param_bytes) + tvm.runtime.load_param_dict(param_bytes) """ - args = [] - for k, v in params.items(): - args.append(k) - args.append(tvm.nd.array(v)) - return _save_param_dict(*args) + return tvm.runtime.save_param_dict(params) def load_param_dict(param_bytes): """Load parameter dictionary to binary bytes. + .. deprecated:: 0.9.0 + Use :py:func:`tvm.runtime.load_param_dict` instead. + Parameters ---------- param_bytes: bytearray @@ -72,7 +69,4 @@ def load_param_dict(param_bytes): params : dict of str to NDArray The parameter dictionary. """ - if isinstance(param_bytes, (bytes, str)): - param_bytes = bytearray(param_bytes) - load_arr = _load_param_dict(param_bytes) - return {v.name: v.array for v in load_arr} + return tvm.runtime.load_param_dict(param_bytes) diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py index 21c06c517bd7..7d58af70afe1 100644 --- a/python/tvm/runtime/__init__.py +++ b/python/tvm/runtime/__init__.py @@ -29,3 +29,4 @@ from .ndarray import vpi, rocm, ext_dev, micro_dev from .module import load_module, enabled, system_lib from .container import String +from .params import save_param_dict, load_param_dict diff --git a/python/tvm/runtime/params.py b/python/tvm/runtime/params.py new file mode 100644 index 000000000000..78e745686c95 --- /dev/null +++ b/python/tvm/runtime/params.py @@ -0,0 +1,69 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. +# pylint: disable=invalid-name +"""Helper utility to save and load parameter dicts.""" +from . import _ffi_api, ndarray + + +def save_param_dict(params): + """Save parameter dictionary to binary bytes. + + The result binary bytes can be loaded by the + GraphModule with API "load_params". + + Parameters + ---------- + params : dict of str to NDArray + The parameter dictionary. + + Returns + ------- + param_bytes: bytearray + Serialized parameters. + + Examples + -------- + .. code-block:: python + + # set up the parameter dict + params = {"param0": arr0, "param1": arr1} + # save the parameters as byte array + param_bytes = tvm.runtime.save_param_dict(params) + # We can serialize the param_bytes and load it back later. + # Pass in byte array to module to directly set parameters + tvm.runtime.load_param_dict(param_bytes) + """ + transformed = {k: ndarray.array(v) for (k, v) in params.items()} + return _ffi_api.SaveParams(transformed) + + +def load_param_dict(param_bytes): + """Load parameter dictionary to binary bytes. + + Parameters + ---------- + param_bytes: bytearray + Serialized parameters. + + Returns + ------- + params : dict of str to NDArray + The parameter dictionary. + """ + if isinstance(param_bytes, (bytes, str)): + param_bytes = bytearray(param_bytes) + return _ffi_api.LoadParams(param_bytes) diff --git a/rust/tvm-graph-rt/src/graph.rs b/rust/tvm-graph-rt/src/graph.rs index 646a20daaf5b..83fe37ea7970 100644 --- a/rust/tvm-graph-rt/src/graph.rs +++ b/rust/tvm-graph-rt/src/graph.rs @@ -483,7 +483,7 @@ named! { ) } -/// Loads a param dict saved using `relay.save_param_dict`. +/// Loads a param dict saved using `runtime.save_param_dict`. pub fn load_param_dict(bytes: &[u8]) -> Result, GraphFormatError> { match parse_param_dict(bytes) { Ok((remaining_bytes, param_dict)) => { diff --git a/rust/tvm-graph-rt/tests/build_model.py b/rust/tvm-graph-rt/tests/build_model.py index d34b4403c936..969075929a42 100755 --- a/rust/tvm-graph-rt/tests/build_model.py +++ b/rust/tvm-graph-rt/tests/build_model.py @@ -23,7 +23,7 @@ import numpy as np import tvm from tvm import te -from tvm import relay +from tvm import relay, runtime from tvm.relay import testing CWD = osp.dirname(osp.abspath(osp.expanduser(__file__))) @@ -47,7 +47,7 @@ def main(): with open(osp.join(CWD, "graph.json"), "w") as f_resnet: f_resnet.write(graph) with open(osp.join(CWD, "graph.params"), "wb") as f_params: - f_params.write(relay.save_param_dict(params)) + f_params.write(runtime.save_param_dict(params)) if __name__ == "__main__": diff --git a/rust/tvm-graph-rt/tests/test_nn/src/build_test_graph.py b/rust/tvm-graph-rt/tests/test_nn/src/build_test_graph.py index e743e48b01f8..0045b3b0557d 100755 --- a/rust/tvm-graph-rt/tests/test_nn/src/build_test_graph.py +++ b/rust/tvm-graph-rt/tests/test_nn/src/build_test_graph.py @@ -23,7 +23,7 @@ import numpy as np import tvm -from tvm import te +from tvm import te, runtime from tvm import relay from tvm.relay import testing @@ -49,7 +49,7 @@ def main(): f_resnet.write(graph) with open(osp.join(out_dir, "graph.params"), "wb") as f_params: - f_params.write(relay.save_param_dict(params)) + f_params.write(runtime.save_param_dict(params)) if __name__ == "__main__": diff --git a/rust/tvm/examples/resnet/src/build_resnet.py b/rust/tvm/examples/resnet/src/build_resnet.py index 03ac611a191a..fdacb5bb1fca 100644 --- a/rust/tvm/examples/resnet/src/build_resnet.py +++ b/rust/tvm/examples/resnet/src/build_resnet.py @@ -27,7 +27,7 @@ import tvm from tvm import te -from tvm import relay +from tvm import relay, runtime from tvm.relay import testing from tvm.contrib import graph_runtime, cc from PIL import Image @@ -88,7 +88,7 @@ def build(target_dir): fo.write(graph) with open(osp.join(target_dir, "deploy_param.params"), "wb") as fo: - fo.write(relay.save_param_dict(params)) + fo.write(runtime.save_param_dict(params)) def download_img_labels(): diff --git a/src/relay/backend/param_dict.cc b/src/relay/backend/param_dict.cc index 1d7e08abcdde..bb0fad9142c1 100644 --- a/src/relay/backend/param_dict.cc +++ b/src/relay/backend/param_dict.cc @@ -31,70 +31,24 @@ #include #include +#include "../../runtime/file_utils.h" + namespace tvm { namespace relay { using namespace runtime; -TVM_REGISTER_GLOBAL("tvm.relay._save_param_dict").set_body([](TVMArgs args, TVMRetValue* rv) { - ICHECK_EQ(args.size() % 2, 0u); - // `args` is in the form "key, value, key, value, ..." - size_t num_params = args.size() / 2; - std::vector names; - names.reserve(num_params); - std::vector arrays; - arrays.reserve(num_params); - for (size_t i = 0; i < num_params * 2; i += 2) { - names.emplace_back(args[i].operator String()); - arrays.emplace_back(args[i + 1].operator DLTensor*()); - } - std::string bytes; - dmlc::MemoryStringStream strm(&bytes); - dmlc::Stream* fo = &strm; - uint64_t header = kTVMNDArrayListMagic, reserved = 0; - fo->Write(header); - fo->Write(reserved); - fo->Write(names); - { - uint64_t sz = static_cast(arrays.size()); - fo->Write(sz); - for (size_t i = 0; i < sz; ++i) { - tvm::runtime::SaveDLTensor(fo, arrays[i]); - } - } - TVMByteArray arr; - arr.data = bytes.c_str(); - arr.size = bytes.length(); - *rv = arr; -}); - -TVM_REGISTER_GLOBAL("tvm.relay._load_param_dict").set_body([](TVMArgs args, TVMRetValue* rv) { - std::string bytes = args[0]; - std::vector names; - dmlc::MemoryStringStream memstrm(&bytes); - dmlc::Stream* strm = &memstrm; - uint64_t header, reserved; - ICHECK(strm->Read(&header)) << "Invalid parameters file format"; - ICHECK(header == kTVMNDArrayListMagic) << "Invalid parameters file format"; - ICHECK(strm->Read(&reserved)) << "Invalid parameters file format"; - ICHECK(strm->Read(&names)) << "Invalid parameters file format"; - uint64_t sz; - strm->Read(&sz, sizeof(sz)); - size_t size = static_cast(sz); - ICHECK(size == names.size()) << "Invalid parameters file format"; - tvm::Array ret; - for (size_t i = 0; i < size; ++i) { - tvm::runtime::NDArray temp; - temp.Load(strm); - auto n = tvm::make_object(); - n->name = std::move(names[i]); - n->array = temp; - ret.push_back(NamedNDArray(n)); - } - *rv = ret; +TVM_REGISTER_GLOBAL("tvm.relay._save_param_dict") + .set_body_typed([](const Map& params) { + std::string s = ::tvm::runtime::SaveParams(params); + // copy return array so it is owned by the ret value + TVMRetValue rv; + rv = TVMByteArray{s.data(), s.size()}; + return rv; + }); +TVM_REGISTER_GLOBAL("tvm.relay._load_param_dict").set_body_typed([](const String& s) { + return ::tvm::runtime::LoadParams(s); }); -TVM_REGISTER_NODE_TYPE(NamedNDArrayNode); - } // namespace relay } // namespace tvm diff --git a/src/relay/backend/param_dict.h b/src/relay/backend/param_dict.h index 384201f94648..96e17a9da07b 100644 --- a/src/relay/backend/param_dict.h +++ b/src/relay/backend/param_dict.h @@ -32,32 +32,7 @@ #include namespace tvm { -namespace relay { - -/*! \brief Magic number for NDArray list file */ -constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; - -/*! - * \brief Wrapper node for naming `NDArray`s. - */ -struct NamedNDArrayNode : public ::tvm::Object { - std::string name; - tvm::runtime::NDArray array; - - void VisitAttrs(tvm::AttrVisitor* v) { - v->Visit("name", &name); - v->Visit("array", &array); - } - - static constexpr const char* _type_key = "NamedNDArray"; - TVM_DECLARE_FINAL_OBJECT_INFO(NamedNDArrayNode, Object); -}; - -class NamedNDArray : public ObjectRef { - public: - TVM_DEFINE_OBJECT_REF_METHODS(NamedNDArray, ObjectRef, NamedNDArrayNode); -}; -} // namespace relay +namespace relay {} // namespace relay } // namespace tvm #endif // TVM_RELAY_BACKEND_PARAM_DICT_H_ diff --git a/src/runtime/file_utils.cc b/src/runtime/file_utils.cc index 3957505a7c7d..92c398b559d2 100644 --- a/src/runtime/file_utils.cc +++ b/src/runtime/file_utils.cc @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -158,5 +159,71 @@ void LoadMetaDataFromFile(const std::string& file_name, void RemoveFile(const std::string& file_name) { std::remove(file_name.c_str()); } +Map LoadParams(const std::string& param_blob) { + dmlc::MemoryStringStream strm(const_cast(¶m_blob)); + return LoadParams(&strm); +} +Map LoadParams(dmlc::Stream* strm) { + Map params; + uint64_t header, reserved; + ICHECK(strm->Read(&header)) << "Invalid parameters file format"; + ICHECK(header == kTVMNDArrayListMagic) << "Invalid parameters file format"; + ICHECK(strm->Read(&reserved)) << "Invalid parameters file format"; + + std::vector names; + ICHECK(strm->Read(&names)) << "Invalid parameters file format"; + uint64_t sz; + strm->Read(&sz); + size_t size = static_cast(sz); + ICHECK(size == names.size()) << "Invalid parameters file format"; + for (size_t i = 0; i < size; ++i) { + // The data_entry is allocated on device, NDArray.load always load the array into CPU. + NDArray temp; + temp.Load(strm); + params.Set(names[i], temp); + } + return params; +} + +void SaveParams(dmlc::Stream* strm, const Map& params) { + std::vector names; + std::vector arrays; + for (auto& p : params) { + names.push_back(p.first); + arrays.push_back(p.second.operator->()); + } + + uint64_t header = kTVMNDArrayListMagic, reserved = 0; + strm->Write(header); + strm->Write(reserved); + strm->Write(names); + { + uint64_t sz = static_cast(arrays.size()); + strm->Write(sz); + for (size_t i = 0; i < sz; ++i) { + tvm::runtime::SaveDLTensor(strm, arrays[i]); + } + } +} + +std::string SaveParams(const Map& params) { + std::string bytes; + dmlc::MemoryStringStream strm(&bytes); + dmlc::Stream* fo = &strm; + SaveParams(fo, params); + return bytes; +} + +TVM_REGISTER_GLOBAL("runtime.SaveParams").set_body_typed([](const Map& params) { + std::string s = ::tvm::runtime::SaveParams(params); + // copy return array so it is owned by the ret value + TVMRetValue rv; + rv = TVMByteArray{s.data(), s.size()}; + return rv; +}); +TVM_REGISTER_GLOBAL("runtime.LoadParams").set_body_typed([](const String& s) { + return ::tvm::runtime::LoadParams(s); +}); + } // namespace runtime } // namespace tvm diff --git a/src/runtime/file_utils.h b/src/runtime/file_utils.h index dfa7d67f1bfe..718d10d5df70 100644 --- a/src/runtime/file_utils.h +++ b/src/runtime/file_utils.h @@ -94,6 +94,32 @@ void LoadMetaDataFromFile(const std::string& file_name, * \param file_name The file name. */ void RemoveFile(const std::string& file_name); + +constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; +/*! + * \brief Load parameters from a string. + * \param param_blob Serialized string of parameters. + * \return Map of parameter name to parameter value. + */ +Map LoadParams(const std::string& param_blob); +/*! + * \brief Load parameters from a stream. + * \param strm Stream to load parameters from. + * \return Map of parameter name to parameter value. + */ +Map LoadParams(dmlc::Stream* strm); +/*! + * \brief Serialize parameters to a byte array. + * \param params Parameters to save. + * \return String containing binary parameter data. + */ +std::string SaveParams(const Map& params); +/*! + * \brief Serialize parameters to a stream. + * \param strm Stream to write to. + * \param params Parameters to save. + */ +void SaveParams(dmlc::Stream* strm, const Map& params); } // namespace runtime } // namespace tvm #endif // TVM_RUNTIME_FILE_UTILS_H_ diff --git a/src/runtime/graph/graph_runtime.cc b/src/runtime/graph/graph_runtime.cc index 6d586cfdd042..6c51e711aef1 100644 --- a/src/runtime/graph/graph_runtime.cc +++ b/src/runtime/graph/graph_runtime.cc @@ -38,6 +38,8 @@ #include #include +#include "../file_utils.h" + namespace tvm { namespace runtime { namespace details { @@ -196,31 +198,10 @@ void GraphRuntime::LoadParams(const std::string& param_blob) { } void GraphRuntime::LoadParams(dmlc::Stream* strm) { - uint64_t header, reserved; - ICHECK(strm->Read(&header)) << "Invalid parameters file format"; - ICHECK(header == kTVMNDArrayListMagic) << "Invalid parameters file format"; - ICHECK(strm->Read(&reserved)) << "Invalid parameters file format"; - - std::vector names; - ICHECK(strm->Read(&names)) << "Invalid parameters file format"; - uint64_t sz; - strm->Read(&sz); - size_t size = static_cast(sz); - ICHECK(size == names.size()) << "Invalid parameters file format"; - for (size_t i = 0; i < size; ++i) { - int in_idx = GetInputIndex(names[i]); - if (in_idx < 0) { - NDArray temp; - temp.Load(strm); - continue; - } - uint32_t eid = this->entry_id(input_nodes_[in_idx], 0); - ICHECK_LT(eid, data_entry_.size()); - - // The data_entry is allocated on device, NDArray.load always load the array into CPU. - NDArray temp; - temp.Load(strm); - data_entry_[eid].CopyFrom(temp); + Map params = ::tvm::runtime::LoadParams(strm); + for (auto& p : params) { + uint32_t eid = this->entry_id(input_nodes_[GetInputIndex(p.first)], 0); + data_entry_[eid].CopyFrom(p.second); } } diff --git a/src/runtime/graph/graph_runtime.h b/src/runtime/graph/graph_runtime.h index 627911883dfb..a1e2ee3b5d74 100644 --- a/src/runtime/graph/graph_runtime.h +++ b/src/runtime/graph/graph_runtime.h @@ -47,9 +47,6 @@ namespace runtime { ICHECK_EQ(ret, 0) << TVMGetLastError(); \ } -/*! \brief Magic number for NDArray list file */ -constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; - /*! \brief operator attributes about tvm op */ struct TVMOpParam { std::string func_name; diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc index 3f890baf52c0..6d121aa67733 100644 --- a/src/runtime/vm/vm.cc +++ b/src/runtime/vm/vm.cc @@ -35,6 +35,8 @@ #include #include +#include "../file_utils.h" + using namespace tvm::runtime; namespace tvm { diff --git a/tests/python/contrib/test_tensorrt.py b/tests/python/contrib/test_tensorrt.py index 60d6b2aa7571..ae8214d6463c 100644 --- a/tests/python/contrib/test_tensorrt.py +++ b/tests/python/contrib/test_tensorrt.py @@ -22,7 +22,7 @@ import tvm import tvm.relay.testing -from tvm import relay +from tvm import relay, runtime from tvm.relay.op.contrib import tensorrt from tvm.contrib import graph_runtime, utils from tvm.runtime.vm import VirtualMachine @@ -265,7 +265,7 @@ def test_tensorrt_serialize_graph_runtime(): def compile_graph(mod, params): with tvm.transform.PassContext(opt_level=3, config={"relay.ext.tensorrt.options": config}): graph, lib, params = relay.build(mod, params=params, target="cuda") - params = relay.save_param_dict(params) + params = runtime.save_param_dict(params) return graph, lib, params def run_graph(graph, lib, params): diff --git a/tests/python/relay/test_cpp_build_module.py b/tests/python/relay/test_cpp_build_module.py index 67f0621ef273..60f3dfa76e38 100644 --- a/tests/python/relay/test_cpp_build_module.py +++ b/tests/python/relay/test_cpp_build_module.py @@ -18,7 +18,7 @@ import tvm from tvm import te -from tvm import relay +from tvm import relay, runtime from tvm.contrib.nvcc import have_fp16 import tvm.testing @@ -86,7 +86,7 @@ def test_fp16_build(): # test rt = tvm.contrib.graph_runtime.create(g_json, mmod, ctx) - rt.load_params(relay.save_param_dict(params)) + rt.load_params(runtime.save_param_dict(params)) rt.run() out = rt.get_output(0) diff --git a/tests/python/relay/test_param_dict.py b/tests/python/relay/test_param_dict.py index 74c9ebcaa355..29e0b5c0463b 100644 --- a/tests/python/relay/test_param_dict.py +++ b/tests/python/relay/test_param_dict.py @@ -17,7 +17,7 @@ import os import numpy as np import tvm -from tvm import te +from tvm import te, runtime import json import base64 from tvm._ffi.base import py_str @@ -31,7 +31,7 @@ def test_save_load(): x = np.ones((10, 2)).astype("float32") y = np.ones((1, 2, 3)).astype("float32") params = {"x": x, "y": y} - param_bytes = relay.save_param_dict(params) + param_bytes = runtime.save_param_dict(params) assert isinstance(param_bytes, bytearray) param2 = relay.load_param_dict(param_bytes) assert len(param2) == 2 @@ -46,7 +46,7 @@ def test_ndarray_reflection(): param_dict = {"x": tvm_array, "y": tvm_array} assert param_dict["x"].same_as(param_dict["y"]) # Serialize then deserialize `param_dict`. - deser_param_dict = relay.load_param_dict(relay.save_param_dict(param_dict)) + deser_param_dict = relay.load_param_dict(runtime.save_param_dict(param_dict)) # Make sure the data matches the original data and `x` and `y` contain the same data. np.testing.assert_equal(deser_param_dict["x"].asnumpy(), tvm_array.asnumpy()) # Make sure `x` and `y` contain the same data. @@ -77,7 +77,7 @@ def verify_graph_runtime(remote, target, shape, dtype): lib = remote.load_module("dev_lib.o") ctx = remote.cpu(0) mod = graph_runtime.create(graph, lib, ctx) - mod.load_params(relay.save_param_dict(params)) + mod.load_params(runtime.save_param_dict(params)) mod.run() out = mod.get_output(0, tvm.nd.empty(shape, dtype=dtype, ctx=ctx)) tvm.testing.assert_allclose(x_in + 1, out.asnumpy()) diff --git a/tests/python/unittest/test_runtime_graph.py b/tests/python/unittest/test_runtime_graph.py index c43a35924420..16e9db42cba3 100644 --- a/tests/python/unittest/test_runtime_graph.py +++ b/tests/python/unittest/test_runtime_graph.py @@ -16,7 +16,7 @@ # under the License. import tvm import tvm.testing -from tvm import te +from tvm import te, runtime import numpy as np import json from tvm import rpc @@ -94,12 +94,12 @@ def check_sharing(): graph, lib, params = relay.build(func, target="llvm", params=params) mod_shared = graph_runtime.create(graph, lib, tvm.cpu(0)) - mod_shared.load_params(relay.save_param_dict(params)) + mod_shared.load_params(runtime.save_param_dict(params)) num_mods = 10 mods = [graph_runtime.create(graph, lib, tvm.cpu(0)) for _ in range(num_mods)] for mod in mods: - mod.share_params(mod_shared, relay.save_param_dict(params)) + mod.share_params(mod_shared, runtime.save_param_dict(params)) a = np.random.uniform(size=(1, 10)).astype("float32") for mod in mods: diff --git a/tests/python/unittest/test_runtime_module_based_interface.py b/tests/python/unittest/test_runtime_module_based_interface.py index 51a587242ae3..a34fe4a062cb 100644 --- a/tests/python/unittest/test_runtime_module_based_interface.py +++ b/tests/python/unittest/test_runtime_module_based_interface.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. import numpy as np -from tvm import relay +from tvm import relay, runtime from tvm.relay import testing import tvm from tvm.contrib import graph_runtime @@ -314,7 +314,7 @@ def verify_cpu_remove_package_params(obj_format): complied_graph_lib_no_params = complied_graph_lib["remove_params"]() complied_graph_lib_no_params.export_library(path_lib) with open(temp.relpath("deploy_param.params"), "wb") as fo: - fo.write(relay.save_param_dict(complied_graph_lib.get_params())) + fo.write(runtime.save_param_dict(complied_graph_lib.get_params())) loaded_lib = tvm.runtime.load_module(path_lib) data = np.random.uniform(-1, 1, size=input_shape(mod)).astype("float32") ctx = tvm.cpu(0) @@ -361,7 +361,7 @@ def verify_gpu_remove_package_params(obj_format): complied_graph_lib_no_params = complied_graph_lib["remove_params"]() complied_graph_lib_no_params.export_library(path_lib) with open(temp.relpath("deploy_param.params"), "wb") as fo: - fo.write(relay.save_param_dict(complied_graph_lib.get_params())) + fo.write(runtime.save_param_dict(complied_graph_lib.get_params())) loaded_lib = tvm.runtime.load_module(path_lib) data = np.random.uniform(-1, 1, size=input_shape(mod)).astype("float32") ctx = tvm.gpu(0) @@ -409,7 +409,7 @@ def verify_rpc_cpu_remove_package_params(obj_format): complied_graph_lib_no_params.export_library(path_lib) path_params = temp.relpath("deploy_param.params") with open(path_params, "wb") as fo: - fo.write(relay.save_param_dict(complied_graph_lib.get_params())) + fo.write(runtime.save_param_dict(complied_graph_lib.get_params())) from tvm import rpc @@ -462,7 +462,7 @@ def verify_rpc_gpu_remove_package_params(obj_format): complied_graph_lib_no_params.export_library(path_lib) path_params = temp.relpath("deploy_param.params") with open(path_params, "wb") as fo: - fo.write(relay.save_param_dict(complied_graph_lib.get_params())) + fo.write(runtime.save_param_dict(complied_graph_lib.get_params())) from tvm import rpc diff --git a/tutorials/frontend/deploy_sparse.py b/tutorials/frontend/deploy_sparse.py index 9641fb8fd14c..98004a93c74f 100644 --- a/tutorials/frontend/deploy_sparse.py +++ b/tutorials/frontend/deploy_sparse.py @@ -81,7 +81,7 @@ import itertools import numpy as np import tensorflow as tf -from tvm import relay +from tvm import relay, runtime from tvm.contrib import graph_runtime from tvm.relay import data_dep_optimization as ddo from tensorflow.python.framework.convert_to_constants import ( @@ -196,7 +196,7 @@ def import_graphdef( with open(os.path.join(abs_path, relay_file), "w") as fo: fo.write(tvm.ir.save_json(mod)) with open(os.path.join(abs_path, relay_params), "wb") as fo: - fo.write(relay.save_param_dict(params)) + fo.write(runtime.save_param_dict(params)) return mod, params, shape_dict