Merge branch 'develop' of https://github.com/PaddlePaddle/GraphNet into develop

Author: JewelRoam

graph_net/torch/blade_disc_backend.py …_net/torch/backend/blade_disc_backend.pygraph_net/torch/blade_disc_backend.py renamed to graph_net/torch/backend/blade_disc_backend.py graph_net/torch/blade_disc_backend.py renamed to graph_net/torch/backend/blade_disc_backend.py

@@ -34,3 +34,6 @@ def __call__(self, model):
     def synchronize(self):
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+
+    def version(self):
+        return torch_blade.version

graph_net/torch/graph_compiler_backend.py …/torch/backend/graph_compiler_backend.pygraph_net/torch/graph_compiler_backend.py renamed to graph_net/torch/backend/graph_compiler_backend.py graph_net/torch/graph_compiler_backend.py renamed to graph_net/torch/backend/graph_compiler_backend.py

@@ -1,12 +1,11 @@
 import torch
+from .graph_compiler_backend import GraphCompilerBackend
 
 try:
     import torch_tensorrt
 except ImportError:
     torch_tensorrt = None
 
-from .graph_compiler_backend import GraphCompilerBackend
-
 
 class TensorRTBackend(GraphCompilerBackend):
     def __call__(self, model):
@@ -16,3 +15,6 @@ def __call__(self, model):
 
     def synchronize(self):
         torch.cuda.synchronize()
+
+    def version(self):
+        return torch_tensorrt.version

graph_net/torch/inductor_backend.py …ph_net/torch/backend/inductor_backend.pygraph_net/torch/inductor_backend.py renamed to graph_net/torch/backend/inductor_backend.py graph_net/torch/inductor_backend.py renamed to graph_net/torch/backend/inductor_backend.py

@@ -14,10 +14,10 @@
 import json
 import numpy as np
 import platform
-from graph_net.torch.graph_compiler_backend import GraphCompilerBackend
-from graph_net.torch.inductor_backend import InductorBackend
-from graph_net.torch.tensorrt_backend import TensorRTBackend
-from graph_net.torch.blade_disc_backend import BladeDISCBackend
+from graph_net.torch.backend.graph_compiler_backend import GraphCompilerBackend
+from graph_net.torch.backend.inductor_backend import InductorBackend
+from graph_net.torch.backend.tensorrt_backend import TensorRTBackend
+from graph_net.torch.backend.blade_disc_backend import BladeDISCBackend
 
 registry_backend = {
     "inductor": InductorBackend(),
@@ -35,6 +35,7 @@ def load_class_from_file(
 
     with open(file_path, "r", encoding="utf-8") as f:
         model_code = f.read()
+    model_code = utils.update_device(model_code, args.device)
     spec = importlib.util.spec_from_loader(module_name, loader=None)
     module = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = module
@@ -196,11 +197,11 @@ def test_single_model(args):
     elif args.compiler == "tensorrt":
         result_data["configuration"][
             "compile_framework_version"
-        ] = f"TensorRT {torch_tensorrt.version}"
+        ] = f"TensorRT {compiler.version}"
     elif args.compiler == "bladedisc":
         result_data["configuration"][
             "compile_framework_version"
-        ] = f"BladeDISC {torch_blade.version}"
+        ] = f"BladeDISC {compiler.version}"
     else:
         result_data["configuration"]["compiler_version"] = "unknown"
 

graph_net/torch/tensorrt_backend.py …ph_net/torch/backend/tensorrt_backend.pygraph_net/torch/tensorrt_backend.py renamed to graph_net/torch/backend/tensorrt_backend.py graph_net/torch/tensorrt_backend.py renamed to graph_net/torch/backend/tensorrt_backend.py

@@ -270,3 +270,16 @@ def replay_tensor(info):
     if dtype is torch.bool:
         return (torch.randn(size=shape) > 0.5).to(dtype).to(device)
     return torch.randn(size=shape).to(dtype).to(device) * std * 0.2 + mean
+
+
+def update_device(code, device):
+    if device == "cuda":
+        pattern = r'device\(type="cpu"\)'
+        replacement = f'device(type="cuda", index={torch.cuda.current_device()})'
+        updated_code = re.sub(pattern, replacement, code)
+        return updated_code
+    else:
+        pattern = r'device\(type="cuda"(?:, index=\d+)?\)'
+        replacement = 'device(type="cpu")'
+        updated_code = re.sub(pattern, replacement, code)
+        return updated_code