add GPU number and lazy load img to GPU

Author: sword4869

arguments/__init__.py

@@ -52,6 +52,7 @@ def __init__(self, parser, sentinel=False):
         self._images = "images"
         self._resolution = -1
         self._white_background = False
+        self._device_number = 0
         self.data_device = "cuda"
         self.eval = False
         super().__init__(parser, "Loading Parameters", sentinel)

render.py

@@ -61,6 +61,6 @@ def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParam
     print("Rendering " + args.model_path)
 
     # Initialize system state (RNG)
-    safe_state(args.quiet)
+    safe_state(args.quiet, args.device_number)
 
     render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test)

scene/cameras.py

@@ -36,14 +36,14 @@ def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
             print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )
             self.data_device = torch.device("cuda")
 
-        self.original_image = image.clamp(0.0, 1.0).to(self.data_device)
+        self.original_image = image.clamp(0.0, 1.0)
         self.image_width = self.original_image.shape[2]
         self.image_height = self.original_image.shape[1]
 
         if gt_alpha_mask is not None:
-            self.original_image *= gt_alpha_mask.to(self.data_device)
+            self.original_image *= gt_alpha_mask
         else:
-            self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device)
+            self.original_image *= torch.ones((1, self.image_height, self.image_width))
 
         self.zfar = 100.0
         self.znear = 0.01

train.py

@@ -211,7 +211,7 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i
     print("Optimizing " + args.model_path)
 
     # Initialize system state (RNG)
-    safe_state(args.quiet)
+    safe_state(args.quiet, args.device_number)
 
     # Start GUI server, configure and run training
     network_gui.init(args.ip, args.port)

utils/general_utils.py

@@ -109,7 +109,7 @@ def build_scaling_rotation(s, r):
     L = R @ L
     return L
 
-def safe_state(silent):
+def safe_state(silent, device_number):
     old_f = sys.stdout
     class F:
         def __init__(self, silent):
@@ -130,4 +130,4 @@ def flush(self):
     random.seed(0)
     np.random.seed(0)
     torch.manual_seed(0)
-    torch.cuda.set_device(torch.device("cuda:0"))
+    torch.cuda.set_device(torch.device(f"cuda:{device_number}"))