Allow to specify device for SubgraphGenerator.

Author: Xreki

graph_net/tools/generate_subgraph_dataset.sh

@@ -204,6 +204,7 @@ EOF
         "model_path_prefix": "${DEVICE_REWRITED_OUTPUT_DIR}",
         "output_dir": "$FUSIBLE_SUBGRAPH_SAMPLES_DIR",
         "subgraph_ranges_json_root": "$FUSIBLE_SUBGRAPH_RANGES_DIR",
+        "device": "cuda",
         "resume": ${RESUME}
     }
 }

graph_net/torch/sample_passes/subgraph_generator.py

@@ -30,6 +30,7 @@ def declare_config(
         subgraph_ranges_json_key: str = "subgraph_ranges",
         group_head_and_tail: bool = False,
         chain_style: bool = False,
+        device: str = "auto",
         resume: bool = False,
         limits_handled_models: int = None,
     ):
@@ -63,7 +64,10 @@ def _has_enough_subgraphs(self, rel_model_path, num_subgraphs):
     def resume(self, rel_model_path: str):
         model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
         torch.cuda.empty_cache()
-        module, inputs = get_torch_module_and_inputs(model_path, use_dummy_inputs=False)
+        device = self._choose_device(self.config["device"])
+        module, inputs = get_torch_module_and_inputs(
+            model_path, use_dummy_inputs=False, device=device
+        )
         gm = parse_sole_graph_module(module, inputs)
         torch.cuda.empty_cache()
         subgraph_ranges = self._get_subgraph_ranges(rel_model_path)
@@ -105,6 +109,11 @@ def fn(submodule, seq_no):
 
         return fn
 
+    def _choose_device(self, device) -> str:
+        if device in ["cpu", "cuda"]:
+            return device
+        return "cuda" if torch.cuda.is_available() else "cpu"
+
 
 class NaiveDecomposerExtractorModule(torch.nn.Module):
     def __init__(