From 3a3220c1febb9a8eabdddc7fdbf8e55ec2525b55 Mon Sep 17 00:00:00 2001 From: Lixing Xiao <111164872+tryhiseyyysum@users.noreply.github.com> Date: Fri, 12 Jul 2024 16:43:00 +0800 Subject: [PATCH 1/4] Update loss_utils.py Updated the function ssim(), the optimized version of ssim() is ssim_optimized(), which reduce the computational complexity. After modified, the function create_window() just need to be called once in the main function in train.py, no need to be called in ssim() in every iteration. --- utils/loss_utils.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/utils/loss_utils.py b/utils/loss_utils.py index 9defc23..77a3f10 100644 --- a/utils/loss_utils.py +++ b/utils/loss_utils.py @@ -25,13 +25,14 @@ def gaussian(window_size, sigma): return gauss / gauss.sum() def create_window(window_size, channel): - _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + """Create a 2D Gaussian window.""" + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) #此处高斯窗口的sigma固定为1.5 _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) return window def ssim(img1, img2, window_size=11, size_average=True): - channel = img1.size(-3) + channel = img1.size(-3) #channel=3 window = create_window(window_size, channel) if img1.is_cuda: @@ -62,3 +63,12 @@ def _ssim(img1, img2, window, window_size, channel, size_average=True): else: return ssim_map.mean(1).mean(1).mean(1) +#-----------------modify----------- +def ssim_optimized(img1, img2, window=None, window_size=11, size_average=True): + channel = img1.size(-3) + if window is None: + window = create_window(window_size, channel).to(img1.device).type_as(img1) + if img1.is_cuda: + window = window.cuda(img1.get_device()) + window = window.type_as(img1) + return _ssim(img1, img2, window, window_size, channel, size_average) From c7ea596c5762a9c78df47894cb43b2d7510db833 Mon Sep 17 00:00:00 2001 From: Lixing Xiao <111164872+tryhiseyyysum@users.noreply.github.com> Date: Fri, 12 Jul 2024 16:45:25 +0800 Subject: [PATCH 2/4] Update train.py Updated the main function, call the create_window() function for just once, since the window_size is fixed to 11 and the channel is fixed to 3 according to the source code. --- train.py | 52 ++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 36 insertions(+), 16 deletions(-) diff --git a/train.py b/train.py index 5d819b3..085b689 100644 --- a/train.py +++ b/train.py @@ -11,8 +11,9 @@ import os import torch +import time from random import randint -from utils.loss_utils import l1_loss, ssim +from utils.loss_utils import l1_loss, ssim, ssim_optimized, create_window from gaussian_renderer import render, network_gui import sys from scene import Scene, GaussianModel @@ -29,27 +30,28 @@ except ImportError: TENSORBOARD_FOUND = False def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from): + start_time=time.time() first_iter = 0 - tb_writer = prepare_output_and_logger(dataset) - gaussians = GaussianModel(dataset.sh_degree) - scene = Scene(dataset, gaussians) - gaussians.training_setup(opt) + tb_writer = prepare_output_and_logger(dataset) # Tensorboard writer + gaussians = GaussianModel(dataset.sh_degree) #高斯模型 + scene = Scene(dataset, gaussians) #场景 + gaussians.training_setup(opt) #训练设置 if checkpoint: (model_params, first_iter) = torch.load(checkpoint) gaussians.restore(model_params, opt) - bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] - background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] #背景颜色 + background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") #背景颜色 - iter_start = torch.cuda.Event(enable_timing = True) - iter_end = torch.cuda.Event(enable_timing = True) + iter_start = torch.cuda.Event(enable_timing = True) #开始时间 + iter_end = torch.cuda.Event(enable_timing = True) #结束时间 viewpoint_stack = None ema_loss_for_log = 0.0 progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress") first_iter += 1 for iteration in range(first_iter, opt.iterations + 1): - if network_gui.conn == None: + if network_gui.conn == None: network_gui.try_connect() while network_gui.conn != None: try: @@ -64,9 +66,9 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi except Exception as e: network_gui.conn = None - iter_start.record() + iter_start.record() #记录开始时间 - gaussians.update_learning_rate(iteration) + gaussians.update_learning_rate(iteration) #更新学习率 # Every 1000 its we increase the levels of SH up to a maximum degree if iteration % 1000 == 0: @@ -83,14 +85,22 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi bg = torch.rand((3), device="cuda") if opt.random_background else background - render_pkg = render(viewpoint_cam, gaussians, pipe, bg) + render_pkg = render(viewpoint_cam, gaussians, pipe, bg) #渲染 image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] - + # Loss + # gt_image = viewpoint_cam.original_image.cuda() + # Ll1 = l1_loss(image, gt_image) + # loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + # loss.backward() + + + # ----------------modify------------- gt_image = viewpoint_cam.original_image.cuda() Ll1 = l1_loss(image, gt_image) - loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim_optimized(image, gt_image, window=window)) loss.backward() + #------------------------------------- iter_end.record() @@ -131,13 +141,18 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi print("\n[ITER {}] Saving Checkpoint".format(iteration)) torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth") + end_time = time.time() # 记录训练结束时间 + total_time = end_time - start_time # 计算总时间 + print(f"\nTraining complete. Total training time: {total_time:.2f} seconds.") # 打印总时间 + + def prepare_output_and_logger(args): if not args.model_path: if os.getenv('OAR_JOB_ID'): unique_str=os.getenv('OAR_JOB_ID') else: unique_str = str(uuid.uuid4()) - args.model_path = os.path.join("./output/", unique_str[0:10]) + args.model_path = os.path.join("/mnt/data1/3dgs_modify_output/", unique_str[0:10]) # Set up output folder print("Output folder: {}".format(args.model_path)) @@ -191,6 +206,11 @@ def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_i torch.cuda.empty_cache() if __name__ == "__main__": + #----------------------create window------------------ + window_size=11 + channel=3 + window=create_window(window_size, channel) + #-------------------------------- # Set up command line argument parser parser = ArgumentParser(description="Training script parameters") lp = ModelParams(parser) From b66e1ad13e4331299c0bfca46f2f43c9b1c278bc Mon Sep 17 00:00:00 2001 From: Lixing Xiao <111164872+tryhiseyyysum@users.noreply.github.com> Date: Fri, 12 Jul 2024 16:56:48 +0800 Subject: [PATCH 3/4] Update loss_utils.py Updated the function ssim(), the optimized version of ssim() is ssim_optimized(), which reduce the computational complexity. After modified, the function create_window() just need to be called once in the main function in train.py, no need to be called in ssim() in every iteration. --- utils/loss_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/loss_utils.py b/utils/loss_utils.py index 77a3f10..ebc929a 100644 --- a/utils/loss_utils.py +++ b/utils/loss_utils.py @@ -26,7 +26,7 @@ def gaussian(window_size, sigma): def create_window(window_size, channel): """Create a 2D Gaussian window.""" - _1D_window = gaussian(window_size, 1.5).unsqueeze(1) #此处高斯窗口的sigma固定为1.5 + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) return window @@ -63,7 +63,7 @@ def _ssim(img1, img2, window, window_size, channel, size_average=True): else: return ssim_map.mean(1).mean(1).mean(1) -#-----------------modify----------- +#-----------------modify------------------------------------ def ssim_optimized(img1, img2, window=None, window_size=11, size_average=True): channel = img1.size(-3) if window is None: From 25b71e16d9e9566bc9a4f5dbfe73fa012c805b6c Mon Sep 17 00:00:00 2001 From: Lixing Xiao <111164872+tryhiseyyysum@users.noreply.github.com> Date: Fri, 12 Jul 2024 16:59:36 +0800 Subject: [PATCH 4/4] Update train.py Updated the main function, call the create_window() function for just once, since the window_size is fixed to 11 and the channel is fixed to 3 according to the source code. Updated the calculation of loss function, call the ssim_optimized(), instead of the original ssim(). --- train.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/train.py b/train.py index 085b689..baaa717 100644 --- a/train.py +++ b/train.py @@ -33,18 +33,18 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi start_time=time.time() first_iter = 0 tb_writer = prepare_output_and_logger(dataset) # Tensorboard writer - gaussians = GaussianModel(dataset.sh_degree) #高斯模型 - scene = Scene(dataset, gaussians) #场景 - gaussians.training_setup(opt) #训练设置 + gaussians = GaussianModel(dataset.sh_degree) + scene = Scene(dataset, gaussians) + gaussians.training_setup(opt) if checkpoint: (model_params, first_iter) = torch.load(checkpoint) gaussians.restore(model_params, opt) - bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] #背景颜色 - background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") #背景颜色 + bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] + background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") - iter_start = torch.cuda.Event(enable_timing = True) #开始时间 - iter_end = torch.cuda.Event(enable_timing = True) #结束时间 + iter_start = torch.cuda.Event(enable_timing = True) + iter_end = torch.cuda.Event(enable_timing = True) viewpoint_stack = None ema_loss_for_log = 0.0 @@ -66,9 +66,9 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi except Exception as e: network_gui.conn = None - iter_start.record() #记录开始时间 + iter_start.record() - gaussians.update_learning_rate(iteration) #更新学习率 + gaussians.update_learning_rate(iteration) # Every 1000 its we increase the levels of SH up to a maximum degree if iteration % 1000 == 0: @@ -85,7 +85,7 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi bg = torch.rand((3), device="cuda") if opt.random_background else background - render_pkg = render(viewpoint_cam, gaussians, pipe, bg) #渲染 + render_pkg = render(viewpoint_cam, gaussians, pipe, bg) image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] # Loss @@ -141,9 +141,9 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi print("\n[ITER {}] Saving Checkpoint".format(iteration)) torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth") - end_time = time.time() # 记录训练结束时间 - total_time = end_time - start_time # 计算总时间 - print(f"\nTraining complete. Total training time: {total_time:.2f} seconds.") # 打印总时间 + end_time = time.time() + total_time = end_time - start_time + print(f"\nTraining complete. Total training time: {total_time:.2f} seconds.") def prepare_output_and_logger(args):