From d85c333d36fd8de0039b5e535d290e84a6489173 Mon Sep 17 00:00:00 2001 From: "sen.li" Date: Wed, 19 Jun 2024 15:33:50 +0800 Subject: [PATCH] update getNvpPnnxModel func comments --- tools/pnnx/src/parse/pnnx_graph_parse.cpp | 6 +++--- tools/pnnx/src/parse/pnnx_graph_parse.h | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tools/pnnx/src/parse/pnnx_graph_parse.cpp b/tools/pnnx/src/parse/pnnx_graph_parse.cpp index e1857f5c609..87af8df7567 100644 --- a/tools/pnnx/src/parse/pnnx_graph_parse.cpp +++ b/tools/pnnx/src/parse/pnnx_graph_parse.cpp @@ -65,9 +65,9 @@ bool PnnxGraph::getNvpPnnxModel(const std::string& pt_path, const std::string& s std::strcpy( argv[argc - 2], end_nodes_info.c_str()); //insert extract_model_name - std::string custom_op_py_info = "extract_model_name=" + custom_op_py; - argv[argc - 1] = new char[custom_op_py_info.size() + 1]; - std::strcpy(argv[argc - 2], custom_op_py_info.c_str()); + std::string extract_model_name_info = "extract_model_name=" + extract_model_name; + argv[argc - 1] = new char[extract_model_name_info.size() + 1]; + std::strcpy(argv[argc - 1], extract_model_name_info.c_str()); int result = main(argc, argv); diff --git a/tools/pnnx/src/parse/pnnx_graph_parse.h b/tools/pnnx/src/parse/pnnx_graph_parse.h index c05704a0398..79669c61f65 100644 --- a/tools/pnnx/src/parse/pnnx_graph_parse.h +++ b/tools/pnnx/src/parse/pnnx_graph_parse.h @@ -15,21 +15,24 @@ class PnnxGraph * @brief Get the Nvp Pnnx Model object * * @param pt_path torchscript path + * @param save_dir the dir of save pnnx * @param input_shape input shape of tensor * @param custom_op_path the path of define custom op * @param custom_op_py the py path of define custom op * @param start_nodes the list of start nodes * @param end_nodes the list of end nodes + * @param extract_model_name the name of which to extract * @return true * @return false */ bool getNvpPnnxModel(const std::string& pt_path, \ - const std::string& save_dir, \ + const std::string& save_dir, \ const std::string& input_shape, \ const std::string& custom_op_path, \ - const std::string& custom_op_py, - const std::string& start_nodes = "", - const std::string& end_nodes = ""); + const std::string& custom_op_py,\ + const std::string& start_nodes = "",\ + const std::string& end_nodes = "",\ + const std::string& extract_model_name = "model");