pnnx load gpu torchscript and reset device (#4330)

tools/pnnx/src/CMakeLists.txt

@@ -4,6 +4,7 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 set(pnnx_pass_level0_SRCS
     pass_level0/constant_unpooling.cpp
     pass_level0/inline_block.cpp
+    pass_level0/reset_device.cpp
     pass_level0/shape_inference.cpp
 )
 

tools/pnnx/src/main.cpp

@@ -327,7 +327,7 @@ int main(int argc, char** argv)
 
     try
     {
-        mod = torch::jit::load(ptpath);
+        mod = torch::jit::load(ptpath, (device == "gpu") ? c10::kCUDA : c10::kCPU);
     }
     catch (const c10::Error& e)
     {
@@ -359,7 +359,7 @@ int main(int argc, char** argv)
     fprintf(stderr, "############# pass_level0\n");
 
     std::map<std::string, pnnx::Attribute> foldable_constants;
-    pnnx::pass_level0(mod, g, input_tensors, input_tensors2, module_operators, ptpath, foldable_constants);
+    pnnx::pass_level0(mod, g, input_tensors, input_tensors2, module_operators, ptpath, device, foldable_constants);
 
     //     g->dump();
 

tools/pnnx/src/pass_level0.cpp

@@ -16,19 +16,22 @@
 
 #include "pass_level0/constant_unpooling.h"
 #include "pass_level0/inline_block.h"
+#include "pass_level0/reset_device.h"
 #include "pass_level0/shape_inference.h"
 
 namespace pnnx {
 
-void pass_level0(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& g, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, std::map<std::string, Attribute>& foldable_constants)
+void pass_level0(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& g, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, const std::string& device, std::map<std::string, Attribute>& foldable_constants)
 {
     inline_block(g, module_operators);
 
+    reset_device(g, device);
+
     constant_unpooling(g);
 
     if (!input_tensors.empty())
     {
-        shape_inference(mod, g, input_tensors, input_tensors2, module_operators, ptpath, foldable_constants);
+        shape_inference(mod, g, input_tensors, input_tensors2, module_operators, ptpath, device, foldable_constants);
     }
 }
 

tools/pnnx/src/pass_level0.h

@@ -20,7 +20,7 @@
 
 namespace pnnx {
 
-void pass_level0(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& g, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, std::map<std::string, Attribute>& foldable_constants);
+void pass_level0(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& g, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, const std::string& device, std::map<std::string, Attribute>& foldable_constants);
 
 } // namespace pnnx
 

tools/pnnx/src/pass_level0/reset_device.cpp

@@ -0,0 +1,36 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// 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 "reset_device.h"
+#include "../pass_level1.h"
+
+namespace pnnx {
+
+void reset_device(std::shared_ptr<torch::jit::Graph>& graph, const std::string& device)
+{
+    for (torch::jit::Node* n : graph->nodes())
+    {
+        if (n->kind().toDisplayString() == std::string("aten::to"))
+        {
+            if (n->hasNamedInput("device"))
+            {
+                torch::jit::Node* device_node = n->namedInput("device")->node();
+
+                device_node->s_(torch::jit::attr::value, (device == "gpu") ? "cuda" : "cpu");
+            }
+        }
+    }
+}
+
+} // namespace pnnx

tools/pnnx/src/pass_level0/reset_device.h

@@ -0,0 +1,21 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// 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 <torch/script.h>
+
+namespace pnnx {
+
+void reset_device(std::shared_ptr<torch::jit::Graph>& graph, const std::string& device);
+
+} // namespace pnnx

tools/pnnx/src/pass_level0/shape_inference.cpp

@@ -17,6 +17,7 @@
 
 #include "pass_level0/constant_unpooling.h"
 #include "pass_level0/inline_block.h"
+#include "pass_level0/reset_device.h"
 #include "pass_level0/shape_inference.h"
 
 namespace pnnx {
@@ -77,7 +78,7 @@ static bool value_link_output(const torch::jit::Value* v, const std::vector<torc
     return false;
 }
 
-void shape_inference(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& graph, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, std::map<std::string, Attribute>& foldable_constants)
+void shape_inference(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& graph, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, const std::string& device, std::map<std::string, Attribute>& foldable_constants)
 {
     // collect all intermediate output tensors
     std::vector<std::unordered_set<std::string> > more_value_names;
@@ -150,13 +151,15 @@ void shape_inference(const torch::jit::Module& mod, std::shared_ptr<torch::jit::
 
         // auto mod2 = mod.deepcopy();
 
-        torch::jit::Module mod2 = torch::jit::load(ptpath);
+        torch::jit::Module mod2 = torch::jit::load(ptpath, (device == "gpu") ? c10::kCUDA : c10::kCPU);
         mod2.eval();
 
         auto graph2 = mod2.get_method("forward").graph();
 
         inline_block(graph2, module_operators);
 
+        reset_device(graph2, device);
+
         constant_unpooling(graph2);
 
         std::vector<torch::jit::Value*> values2;

tools/pnnx/src/pass_level0/shape_inference.h

@@ -18,6 +18,6 @@
 
 namespace pnnx {
 
-void shape_inference(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& graph, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, std::map<std::string, Attribute>& foldable_constants);
+void shape_inference(const torch::jit::Module& mod, std::shared_ptr<torch::jit::Graph>& graph, const std::vector<at::Tensor>& input_tensors, const std::vector<at::Tensor>& input_tensors2, const std::vector<std::string>& module_operators, const std::string& ptpath, const std::string& device, std::map<std::string, Attribute>& foldable_constants);
 
 } // namespace pnnx