Map weights to CUDA or CPU depending on which is enabled (#216)

Author: jjaime2

ada_feeding_action_select/ada_feeding_action_select/adapters/hapticnet_adapter.py

@@ -47,10 +47,13 @@ def __init__(
 
         # Init CUDA
         self.use_cuda = torch.cuda.is_available()
+        self.device = torch.device("cuda") if self.use_cuda else torch.device("cpu")
         if self.use_cuda:
             logger.info("Init HapticNet with CUDA")
             os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
             os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_index)
+        else:
+            logger.info("Init HapticNet with CPU")
 
         # Init HapticNet
         self.config = HapticNetConfig(n_output=n_features)
@@ -60,7 +63,7 @@ def __init__(
         ckpt_file = os.path.join(
             get_package_share_directory("ada_feeding_action_select"), "data", checkpoint
         )
-        ckpt = torch.load(ckpt_file)
+        ckpt = torch.load(ckpt_file, map_location=self.device)
         self.hapticnet.load_state_dict(ckpt["state_dict"])
         self.hapticnet.eval()
         if self.use_cuda:

ada_feeding_action_select/ada_feeding_action_select/adapters/spanet_adapter.py

@@ -49,10 +49,13 @@ def __init__(
 
         # Init CUDA
         self.use_cuda = torch.cuda.is_available()
+        self.device = torch.device("cuda") if self.use_cuda else torch.device("cpu")
         if self.use_cuda:
             logger.info("Init SPANet with CUDA")
             os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
             os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_index)
+        else:
+            logger.info("Init SPANet with CPU")
 
         # Init SPANet
         self.config = SPANetConfig(n_features=n_features)
@@ -62,7 +65,7 @@ def __init__(
         ckpt_file = os.path.join(
             get_package_share_directory("ada_feeding_action_select"), "data", checkpoint
         )
-        ckpt = torch.load(ckpt_file)
+        ckpt = torch.load(ckpt_file, map_location=self.device)
         self.spanet.load_state_dict(ckpt["net"])
         self.spanet.eval()
         if self.use_cuda:

ada_feeding_action_select/ada_feeding_action_select/policy_service.py

@@ -226,7 +226,7 @@ def _init_checkpoints_record(self, context_cls: type, posthoc_cls: type) -> None
             )
             if len(pt_files) > 0:
                 with open(pt_files[-1], "rb") as ckpt_file:
-                    ckpt = torch.load(ckpt_file)
+                    ckpt = torch.load(ckpt_file, map_location=self.device)
                     try:
                         if ckpt["context_cls"] != context_cls:
                             self.get_logger().warning(
@@ -253,6 +253,8 @@ def __init__(self):
         register_logger(self.get_logger())
         self._declare_parameters()
 
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         # Name of the Policy
         policy_name = self.get_parameter("policy").value
         if policy_name is None: