From 989320fdf2563dc06bef852dc78ebd2cf820c5ba Mon Sep 17 00:00:00 2001 From: Pythonix Huang <370048201@qq.com> Date: Thu, 13 Jul 2023 02:30:45 +0800 Subject: [PATCH 1/7] Provide --data_device option to put data on CPU to save VRAM for training (#14) * Provide --data_on_cpu option to save VRAM for training when there are many training images such as in large scene, most of the VRAM are used to store training data, use --data_on_cpu can help reduce VRAM and make it possible to train on GPU with less VRAM * Fix data_on_cpu effect on default mask * --data_on_cpu to --data_device * update readme * format warning infos --- README.md | 2 ++ arguments/__init__.py | 1 + scene/cameras.py | 15 +++++++++++---- utils/camera_utils.py | 2 +- 4 files changed, 15 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 08e7587..22e744e 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,8 @@ python train.py -s Space-separated iterations at which the training script saves the Gaussian model, ```7000 30000 ``` by default. #### --quiet Flag to omit any text written to standard out pipe. + #### --data_device + Specify where to put the data on,```cuda``` by default, recommend use ```cpu``` if training on large scale/resolution dataset, will save a lot of VRAM required to train, but slightly slower the training
diff --git a/arguments/__init__.py b/arguments/__init__.py index 4ab86f4..eba1dba 100644 --- a/arguments/__init__.py +++ b/arguments/__init__.py @@ -52,6 +52,7 @@ class ModelParams(ParamGroup): self._images = "images" self._resolution = -1 self._white_background = False + self.data_device = "cuda" self.eval = False super().__init__(parser, "Loading Parameters", sentinel) diff --git a/scene/cameras.py b/scene/cameras.py index b57d351..abf6e52 100644 --- a/scene/cameras.py +++ b/scene/cameras.py @@ -17,7 +17,7 @@ from utils.graphics_utils import getWorld2View2, getProjectionMatrix class Camera(nn.Module): def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, image_name, uid, - trans=np.array([0.0, 0.0, 0.0]), scale=1.0 + trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda" ): super(Camera, self).__init__() @@ -29,14 +29,21 @@ class Camera(nn.Module): self.FoVy = FoVy self.image_name = image_name - self.original_image = image.clamp(0.0, 1.0).cuda() + try: + self.data_device = torch.device(data_device) + except Exception as e: + print(e) + 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.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.cuda() + self.original_image *= gt_alpha_mask.to(self.data_device) else: - self.original_image *= torch.ones((1, self.image_height, self.image_width), device="cuda") + self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device) self.zfar = 100.0 self.znear = 0.01 diff --git a/utils/camera_utils.py b/utils/camera_utils.py index 21d762e..1a54d0a 100644 --- a/utils/camera_utils.py +++ b/utils/camera_utils.py @@ -49,7 +49,7 @@ def loadCam(args, id, cam_info, resolution_scale): return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, FoVx=cam_info.FovX, FoVy=cam_info.FovY, image=gt_image, gt_alpha_mask=loaded_mask, - image_name=cam_info.image_name, uid=id) + image_name=cam_info.image_name, uid=id, data_device=args.data_device) def cameraList_from_camInfos(cam_infos, resolution_scale, args): camera_list = [] From 8abb25b632688dadc9ea0f78d69e78fce9caddf3 Mon Sep 17 00:00:00 2001 From: Snosixtyboo <40643808+Snosixtyboo@users.noreply.github.com> Date: Wed, 12 Jul 2023 21:24:13 +0200 Subject: [PATCH 2/7] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 22e744e..ce84c04 100644 --- a/README.md +++ b/README.md @@ -166,7 +166,7 @@ python train.py -s #### --quiet Flag to omit any text written to standard out pipe. #### --data_device - Specify where to put the data on,```cuda``` by default, recommend use ```cpu``` if training on large scale/resolution dataset, will save a lot of VRAM required to train, but slightly slower the training + Specify where to put the source image data, ```cuda``` by default, recommended to use ```cpu``` if training on large/high-resolution dataset, will reduce VRAM consumption, but slightly slow down training.
From d843dc1224605d684049f9176b866409e3ee7f22 Mon Sep 17 00:00:00 2001 From: Snosixtyboo <40643808+Snosixtyboo@users.noreply.github.com> Date: Wed, 12 Jul 2023 21:25:21 +0200 Subject: [PATCH 3/7] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ce84c04..172382d 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,8 @@ python train.py -s Add this flag to use a MipNeRF360-style training/test split for evaluation. #### --resolution / -r Specifies resolution of the loaded images before training. If provided ```1, 2, 4``` or ```8```, uses original, 1/2, 1/4 or 1/8 resolution, respectively. For all other values, rescales the width to the given number while maintaining image aspect. **If not set and input image width exceeds 1.6K pixels, inputs are automatically rescaled to this target.** + #### --data_device + Specifies where to put the source image data, ```cuda``` by default, recommended to use ```cpu``` if training on large/high-resolution dataset, will reduce VRAM consumption, but slightly slow down training. #### --white_background / -w Add this flag to use white background instead of black (default), e.g., for evaluation of NeRF Synthetic dataset. #### --sh_degree @@ -165,8 +167,6 @@ python train.py -s Space-separated iterations at which the training script saves the Gaussian model, ```7000 30000 ``` by default. #### --quiet Flag to omit any text written to standard out pipe. - #### --data_device - Specify where to put the source image data, ```cuda``` by default, recommended to use ```cpu``` if training on large/high-resolution dataset, will reduce VRAM consumption, but slightly slow down training.
From 770f6b05e2a31b14e2b5fe15665b16f3dca70da2 Mon Sep 17 00:00:00 2001 From: bkerbl Date: Wed, 12 Jul 2023 12:10:03 +0200 Subject: [PATCH 4/7] Clarification --- full_eval.py | 4 ++-- scene/gaussian_model.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/full_eval.py b/full_eval.py index d174ef8..672ee12 100644 --- a/full_eval.py +++ b/full_eval.py @@ -12,7 +12,7 @@ import os from argparse import ArgumentParser -mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"] +mipnerf360_outdoor_scenes = ["flowers", "garden", "stump", "treehill"] mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"] tanks_and_temples_scenes = ["truck", "train"] deep_blending_scenes = ["drjohnson", "playroom"] @@ -37,7 +37,7 @@ if not args.skip_training or not args.skip_rendering: args = parser.parse_args() if not args.skip_training: - common_args = " --quiet --eval --test_iterations -1" + common_args = " --eval --save_iterations -1" for scene in mipnerf360_outdoor_scenes: source = args.mipnerf360 + "/" + scene os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args) diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index 20a81d4..f50a4b2 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -84,7 +84,9 @@ class GaussianModel: self.active_sh_degree += 1 def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): + spatial_lr_scale = 5 self.spatial_lr_scale = spatial_lr_scale + #print(spatial_lr_scale) fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() From 2d2d5ce1c365981a387269d0e5f43cbd5cabfccc Mon Sep 17 00:00:00 2001 From: bkerbl Date: Wed, 12 Jul 2023 12:14:22 +0200 Subject: [PATCH 5/7] Undo spatial lr removal --- full_eval.py | 4 ++-- scene/gaussian_model.py | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/full_eval.py b/full_eval.py index 672ee12..d174ef8 100644 --- a/full_eval.py +++ b/full_eval.py @@ -12,7 +12,7 @@ import os from argparse import ArgumentParser -mipnerf360_outdoor_scenes = ["flowers", "garden", "stump", "treehill"] +mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"] mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"] tanks_and_temples_scenes = ["truck", "train"] deep_blending_scenes = ["drjohnson", "playroom"] @@ -37,7 +37,7 @@ if not args.skip_training or not args.skip_rendering: args = parser.parse_args() if not args.skip_training: - common_args = " --eval --save_iterations -1" + common_args = " --quiet --eval --test_iterations -1" for scene in mipnerf360_outdoor_scenes: source = args.mipnerf360 + "/" + scene os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args) diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index f50a4b2..20a81d4 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -84,9 +84,7 @@ class GaussianModel: self.active_sh_degree += 1 def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): - spatial_lr_scale = 5 self.spatial_lr_scale = spatial_lr_scale - #print(spatial_lr_scale) fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() From 73ad40f414ea7071548ccd1f7e90a60d0cd0cb9d Mon Sep 17 00:00:00 2001 From: bkerbl Date: Wed, 12 Jul 2023 15:38:28 +0200 Subject: [PATCH 6/7] fixed learning rate for spatial --- scene/gaussian_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index 20a81d4..699f02c 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -84,7 +84,7 @@ class GaussianModel: self.active_sh_degree += 1 def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): - self.spatial_lr_scale = spatial_lr_scale + self.spatial_lr_scale = 5#spatial_lr_scale fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() From b1a35bc94bfb1a94c6f8aeca341978b4d32e094b Mon Sep 17 00:00:00 2001 From: bkerbl Date: Thu, 13 Jul 2023 13:12:45 +0200 Subject: [PATCH 7/7] Changed alpha init spatial lr behavior --- scene/gaussian_model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index 699f02c..cd8c2ce 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -84,7 +84,7 @@ class GaussianModel: self.active_sh_degree += 1 def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): - self.spatial_lr_scale = 5#spatial_lr_scale + self.spatial_lr_scale = 5 fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() @@ -98,7 +98,7 @@ class GaussianModel: rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda") rots[:, 0] = 1 - opacities = inverse_sigmoid(0.5 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda")) + opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda")) self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True)) self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True)) @@ -170,7 +170,7 @@ class GaussianModel: PlyData([el]).write(path) def reset_opacity(self): - opacities_new = inverse_sigmoid(torch.ones_like(self.get_opacity)*0.01) + opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity") self._opacity = optimizable_tensors["opacity"]