Crashes on non-GPU machines, or if config.pytorch.gpu == -1 (proposed fix included)

[petrock99]

Crashes on non-GPU machines, or if config.pytorch.gpu == -1

Python 3.6.13 with the same packages as in the EPro-PnP-6PoF README.md

I tried to run this on a mac to try to step through the code with a visual IDE/debugger, like PyCharm. It crashes pretty early. Same if config.pytorch.gpu is -1 on GPU hardware. I had to make a few modifications to get it to work. While I was in there I also added a --forceCPU argument to force it to run on CPU when on GPU hardware without messing with the config files. Its something I like to add to all my projects. I tried making a PR, but GitHub won't let me. Here is the diff (file also attached):

diff.txt

diff --git a/EPro-PnP-6DoF/lib/config.py b/EPro-PnP-6DoF/lib/config.py
index 1d48f5d..1a2bf2a 100644
--- a/EPro-PnP-6DoF/lib/config.py
+++ b/EPro-PnP-6DoF/lib/config.py
@@ -13,6 +13,7 @@ import sys
 import ref
 from datetime import datetime
 import utils.fancy_logger as logger
+import torch
 from tensorboardX import SummaryWriter
 from pprint import pprint

@@ -29,7 +30,7 @@ def get_default_config_pytorch():
     config = edict()
     config.exp_id = 'CDPN'          # Experiment ID
     config.
task = 'rot'             # 'rot | trans | trans_rot'
-    config.gpu = 1
+    config.gpu = 1 if torch.cuda.is_available() else -1
     config.threads_num = 12         # 'nThreads'
     config.debug = False
     config.demo = '../demo'         # demo save path
@@ -177,9 +178,14 @@ class config():
     def __init__(self):
         self.parser = argparse.ArgumentParser(description='pose experiment')
         self.parser.add_argument('--cfg', required=True, type=str, help='path/to/configure_file')
-        self.parser.add_argument('--load_model', type=str, help='path/to/model, requird when resume/test')
+        self.parser.add_argument('--load_model', type=str, help='path/to/model, required when resume/test')
         self.parser.add_argument('--debug', action='store_true', help='')
         self.parser.add_argument('--test', action='store_true', help='')
+        self.parser.add_argument('--force_cpu',
+                                 action='store_true',
+                                 help='Force running on the CPU instead of GPU. Default: Run on GPU if available, otherwise CPU.',
+                                 default=False,
+                                 required=False)

     def parse(self):
         config = get_base_config()                  # get default arguments
@@ -188,6 +194,10 @@ class config():
             config.pytorch[k] = v 
         config_file = config.pytorch.cfg
         config = update_config_from_file(config, config_file, check_necessity=False) # update arguments from config file
+
+        if config.pytorch.force_cpu or not torch.cuda.is_available():
+            config.pytorch.gpu = -1
+
         # complement config regarding dataset
         if config.dataset.name.lower() in ['lm', 'lmo']:
             config.dataset['camera_matrix'] = ref.K
diff --git a/EPro-PnP-6DoF/lib/test.py b/EPro-PnP-6DoF/lib/test.py
index c6ee718..e8406f3 100644
--- a/EPro-PnP-6DoF/lib/test.py
+++ b/EPro-PnP-6DoF/lib/test.py
@@ -86,14 +86,18 @@ def test(epoch, cfg, data_loader, model, obj_vtx, obj_info, criterions):
         os.makedirs(vis_dir)

     cam_intrinsic_np = cfg.dataset.camera_matrix.astype(np.float32)
-    cam_intrinsic = torch.from_numpy(cam_intrinsic_np).cuda(cfg.pytorch.gpu)
+    cam_intrinsic = torch.from_numpy(cam_intrinsic_np)

     epropnp = EProPnP6DoF(
         mc_samples=512,
         num_iter=4,
         solver=LMSolver(
             dof=6,
-            num_iter=3)).cuda(cfg.pytorch.gpu)
+            num_iter=3))
+
+    if cfg.pytorch.gpu > -1:
+        cam_intrinsic = cam_intrinsic.cuda(cfg.pytorch.gpu)
+        epropnp = epropnp.cuda(cfg.pytorch.gpu)

     for i, (obj, obj_id, inp, pose, c_box, s_box, box, trans_local) in enumerate(data_loader):
         if cfg.pytorch.gpu > -1:
diff --git a/EPro-PnP-6DoF/lib/train.py b/EPro-PnP-6DoF/lib/train.py
index 3a41c06..3d2d9d8 100644
--- a/EPro-PnP-6DoF/lib/train.py
+++ b/EPro-PnP-6DoF/lib/train.py
@@ -42,7 +42,7 @@ def train(epoch, cfg, data_loader, model, obj_info, criterions, optimizer=None):
         os.makedirs(vis_dir)

     cam_intrinsic_np = cfg.dataset.camera_matrix.astype(np.float32)
-    cam_intrinsic = torch.from_numpy(cam_intrinsic_np).cuda(cfg.pytorch.gpu)
+    cam_intrinsic = torch.from_numpy(cam_intrinsic_np)

     epropnp = EProPnP6DoF(
         mc_samples=512,
@@ -54,7 +54,11 @@ def train(epoch, cfg, data_loader, model, obj_info, criterions, optimizer=None):
                 dof=6,
                 num_points=16,
                 num_proposals=4,
-                num_iter=3))).cuda(cfg.pytorch.gpu)
+                num_iter=3)))
+
+    if cfg.pytorch.gpu > -1:
+        cam_intrinsic = cam_intrinsic.cuda(cfg.pytorch.gpu)
+        epropnp = epropnp.cuda(cfg.pytorch.gpu)

     for i, (obj, obj_id, inp, target, loss_msk, trans_local, pose, c_box, s_box, box) in enumerate(data_loader):
         cur_iter = i + (epoch - 1) * num_iters