diff --git a/arguments/__init__.py b/arguments/__init__.py index dfcfea9..8ba03db 100644 --- a/arguments/__init__.py +++ b/arguments/__init__.py @@ -134,6 +134,7 @@ class OptimizationParams(ParamGroup): self.opacity_threshold_coarse = 0.005 self.opacity_threshold_fine_init = 0.005 self.opacity_threshold_fine_after = 0.005 + self.batch_size=1, super().__init__(parser, "Optimization Parameters") diff --git a/arguments/dynerf/default.py b/arguments/dynerf/default.py index 57453fa..29035f7 100644 --- a/arguments/dynerf/default.py +++ b/arguments/dynerf/default.py @@ -6,22 +6,24 @@ ModelHiddenParams = dict( 'resolution': [64, 64, 64, 150] }, multires = [1,2,4,8], - defor_depth = 2, + defor_depth = 1, net_width = 256, plane_tv_weight = 0.0002, time_smoothness_weight = 0.001, l1_time_planes = 0.001, + no_do=False ) OptimizationParams = dict( dataloader=True, iterations = 30_000, + batch_size=4, coarse_iterations = 3000, - densify_until_iter = 15_000, - opacity_reset_interval = 6000, + densify_until_iter = 20_000, + opacity_reset_interval = 20000, - opacity_threshold_coarse = 0.005, - opacity_threshold_fine_init = 0.005, - opacity_threshold_fine_after = 0.005, + opacity_threshold_coarse = 0.05, + opacity_threshold_fine_init = 0.05, + opacity_threshold_fine_after = 0.05, # pruning_interval = 2000 ) \ No newline at end of file diff --git a/arguments/dynerf_2/default.py b/arguments/dynerf_2/default.py new file mode 100644 index 0000000..5428474 --- /dev/null +++ b/arguments/dynerf_2/default.py @@ -0,0 +1,29 @@ +ModelHiddenParams = dict( + kplanes_config = { + 'grid_dimensions': 2, + 'input_coordinate_dim': 4, + 'output_coordinate_dim': 16, + 'resolution': [64, 64, 64, 150] + }, + multires = [1,2,4,8], + defor_depth = 1, + net_width = 256, + plane_tv_weight = 0.0002, + time_smoothness_weight = 0.001, + l1_time_planes = 0.001, + no_do=False + +) +OptimizationParams = dict( + dataloader=True, + iterations = 20_000, + batch_size=8, + coarse_iterations = 3000, + densify_until_iter = 20_000, + opacity_reset_interval = 3000, + + opacity_threshold_coarse = 0.05, + opacity_threshold_fine_init = 0.05, + opacity_threshold_fine_after = 0.05, + # pruning_interval = 2000 +) \ No newline at end of file diff --git a/arguments/dynerf_3/default.py b/arguments/dynerf_3/default.py new file mode 100644 index 0000000..99221a4 --- /dev/null +++ b/arguments/dynerf_3/default.py @@ -0,0 +1,29 @@ +ModelHiddenParams = dict( + kplanes_config = { + 'grid_dimensions': 2, + 'input_coordinate_dim': 4, + 'output_coordinate_dim': 16, + 'resolution': [64, 64, 64, 150] + }, + multires = [1,2,4,8], + defor_depth = 1, + net_width = 256, + plane_tv_weight = 0.0002, + time_smoothness_weight = 0.001, + l1_time_planes = 0.001, + no_do=False + +) +OptimizationParams = dict( + dataloader=True, + iterations = 60_000, + batch_size=8, + coarse_iterations = 3000, + densify_until_iter = 20_000, + opacity_reset_interval = 20000, + + opacity_threshold_coarse = 0.05, + opacity_threshold_fine_init = 0.05, + opacity_threshold_fine_after = 0.05, + # pruning_interval = 2000 +) \ No newline at end of file diff --git a/arguments/dynerf_4/default.py b/arguments/dynerf_4/default.py new file mode 100644 index 0000000..bb4c4c6 --- /dev/null +++ b/arguments/dynerf_4/default.py @@ -0,0 +1,29 @@ +ModelHiddenParams = dict( + kplanes_config = { + 'grid_dimensions': 2, + 'input_coordinate_dim': 4, + 'output_coordinate_dim': 16, + 'resolution': [64, 64, 64, 150] + }, + multires = [1,2,4,8], + defor_depth = 1, + net_width = 256, + plane_tv_weight = 0.0002, + time_smoothness_weight = 0.001, + l1_time_planes = 0.001, + no_do=False + +) +OptimizationParams = dict( + dataloader=True, + iterations = 60_000, + batch_size=1, + coarse_iterations = 3000, + densify_until_iter = 40_000, + opacity_reset_interval = 20000, + + opacity_threshold_coarse = 0.05, + opacity_threshold_fine_init = 0.05, + opacity_threshold_fine_after = 0.05, + # pruning_interval = 2000 +) \ No newline at end of file diff --git a/arguments/dynerf_static/default.py b/arguments/dynerf_static/default.py new file mode 100644 index 0000000..3762828 --- /dev/null +++ b/arguments/dynerf_static/default.py @@ -0,0 +1,34 @@ +ModelHiddenParams = dict( + kplanes_config = { + 'grid_dimensions': 2, + 'input_coordinate_dim': 4, + 'output_coordinate_dim': 16, + 'resolution': [64, 64, 64, 150] + }, + multires = [1,2,4,8], + defor_depth = 1, + net_width = 256, + plane_tv_weight = 0.0002, + time_smoothness_weight = 0.001, + l1_time_planes = 0.001, + no_do=False + +) +OptimizationParams = dict( + dataloader=True, + iterations = 10_000, + batch_size=8, + coarse_iterations = 10000, + densify_until_iter = 20_000, + opacity_reset_interval = 3000, + + opacity_threshold_coarse = 0.05, + opacity_threshold_fine_init = 0.05, + opacity_threshold_fine_after = 0.05, + # pruning_interval = 2000 + # deformation_lr_init = 0.00016, + # deformation_lr_final = 0.000016, + # deformation_lr_delay_mult = 0.01, + # grid_lr_init = 0.0016, + # grid_lr_final = 0.00016, +) \ No newline at end of file diff --git a/arguments/hypernerf/default.py b/arguments/hypernerf/default.py index 009e278..7544e38 100644 --- a/arguments/hypernerf/default.py +++ b/arguments/hypernerf/default.py @@ -15,9 +15,10 @@ ModelHiddenParams = dict( ) OptimizationParams = dict( dataloader=False, - iterations = 60_000, + iterations = 30000, + batch_size=1, coarse_iterations = 3000, - densify_until_iter = 45_000, + densify_until_iter = 20_000, opacity_reset_interval = 6000, # position_lr_init = 0.00016, # position_lr_final = 0.0000016, diff --git a/scene/grid.py b/scene/grid.py new file mode 100644 index 0000000..42272ee --- /dev/null +++ b/scene/grid.py @@ -0,0 +1,63 @@ +import os +import time +import functools +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + +parent_dir = os.path.dirname(os.path.abspath(__file__)) + + +def create_grid(type, **kwargs): + if type == 'DenseGrid': + return DenseGrid(**kwargs) + elif type == 'TensoRFGrid': + return TensoRFGrid(**kwargs) + else: + raise NotImplementedError + + +''' Dense 3D grid +''' +class DenseGrid(nn.Module): + def __init__(self, channels, world_size, xyz_min, xyz_max, **kwargs): + super(DenseGrid, self).__init__() + self.channels = channels + self.world_size = world_size + self.register_buffer('xyz_min', torch.Tensor(xyz_min)) + self.register_buffer('xyz_max', torch.Tensor(xyz_max)) + self.grid = nn.Parameter(torch.zeros([1, channels, *world_size])) + + def forward(self, xyz): + ''' + xyz: global coordinates to query + ''' + shape = xyz.shape[:-1] + xyz = xyz.reshape(1,1,1,-1,3) + ind_norm = ((xyz - self.xyz_min) / (self.xyz_max - self.xyz_min)).flip((-1,)) * 2 - 1 + out = F.grid_sample(self.grid, ind_norm, mode='bilinear', align_corners=True) + out = out.reshape(self.channels,-1).T.reshape(*shape,self.channels) + if self.channels == 1: + out = out.squeeze(-1) + return out + + def scale_volume_grid(self, new_world_size): + if self.channels == 0: + self.grid = nn.Parameter(torch.zeros([1, self.channels, *new_world_size])) + else: + self.grid = nn.Parameter( + F.interpolate(self.grid.data, size=tuple(new_world_size), mode='trilinear', align_corners=True)) + + def get_dense_grid(self): + return self.grid + + @torch.no_grad() + def __isub__(self, val): + self.grid.data -= val + return self + + def extra_repr(self): + return f'channels={self.channels}, world_size={self.world_size.tolist()}' + diff --git a/scene/neural_3D_dataset_NDC.py b/scene/neural_3D_dataset_NDC.py index 63bbcad..ae90d5f 100644 --- a/scene/neural_3D_dataset_NDC.py +++ b/scene/neural_3D_dataset_NDC.py @@ -264,6 +264,7 @@ class Neural3D_NDC_Dataset(Dataset): self.near_fars = poses_arr[:, -2:] videos = glob.glob(os.path.join(self.root_dir, "cam*")) videos = sorted(videos) + breakpoint() assert len(videos) == poses_arr.shape[0] H, W, focal = poses[0, :, -1] diff --git a/scripts/train_dynerf.sh b/scripts/train_dynerf.sh index b8c07e2..48cc0b3 100644 --- a/scripts/train_dynerf.sh +++ b/scripts/train_dynerf.sh @@ -1,18 +1,27 @@ exp_name=$1 -export CUDA_VISIBLE_DEVICES=0&&python train.py -s data/hypernerf/misc/split-cookie --port 6068 --expname "$exp_name/split-cookie" --configs arguments/$exp_name/default.py & -export CUDA_VISIBLE_DEVICES=1&&python train.py -s data/hypernerf/virg/vrig-3dprinter --port 6066 --expname "$exp_name/3dprinter" --configs arguments/$exp_name/default.py & -export CUDA_VISIBLE_DEVICES=2&&python train.py -s data/hypernerf/interp/chickchicken --port 6069 --expname "$exp_name/interp-chicken" --configs arguments/$exp_name/default.py & -export CUDA_VISIBLE_DEVICES=3&&python train.py -s data/hypernerf/interp/cut-lemon1 --port 6070 --expname "$exp_name/cut-lemon1" --configs arguments/$exp_name/cut-lemon1.py & -# export CUDA_VISIBLE_DEVICES=3&&python train.py -s data/hypernerf/interp/hand1-dense-v2 --port 6071 --expname "$exp_name/hand1-dense-v2" --configs arguments/$exp_name/hand1-dense-v2.py +export CUDA_VISIBLE_DEVICES=0&&python train.py -s data/dynerf/cut_roasted_beef --port 6068 --expname "$exp_name/cut_roasted_beef" --configs arguments/$exp_name/default.py & +export CUDA_VISIBLE_DEVICES=1&&python train.py -s data/dynerf/cook_spinach --port 6066 --expname "$exp_name/cook_spinach" --configs arguments/$exp_name/default.py & +export CUDA_VISIBLE_DEVICES=2&&python train.py -s data/dynerf/sear_steak --port 6069 --expname "$exp_name/sear_steak" --configs arguments/$exp_name/default.py & wait -export CUDA_VISIBLE_DEVICES=2&&python render.py --model_path output/$exp_name/cut_roasted_beef --configs arguments/$exp_name/default.py --skip_train --skip_test& -export CUDA_VISIBLE_DEVICES=3&&python render.py --model_path output/$exp_name/cook_spinach --configs arguments/$exp_name/default.py --skip_train --skip_test & -export CUDA_VISIBLE_DEVICES=2&&python render.py --model_path output/$exp_name/sear_steak --configs arguments/$exp_name/default.py --skip_train --skip_test& +export CUDA_VISIBLE_DEVICES=0&&python train.py -s data/dynerf/flame_salmon_1 --port 6070 --expname "$exp_name/flame_salmon_1" --configs arguments/$exp_name/default.py & +export CUDA_VISIBLE_DEVICES=1&&python train.py -s data/dynerf/flame_steak --port 6071 --expname "$exp_name/flame_steak" --configs arguments/$exp_name/default.py & +export CUDA_VISIBLE_DEVICES=2&&python train.py -s data/dynerf/coffee_martini --port 6071 --expname "$exp_name/coffee_martini" --configs arguments/$exp_name/default.py & +wait +export CUDA_VISIBLE_DEVICES=0&&python render.py --model_path output/$exp_name/cut_roasted_beef --configs arguments/$exp_name/default.py --skip_train & +export CUDA_VISIBLE_DEVICES=1&&python render.py --model_path output/$exp_name/cook_spinach --configs arguments/$exp_name/default.py --skip_train & +export CUDA_VISIBLE_DEVICES=2&&python render.py --model_path output/$exp_name/sear_steak --configs arguments/$exp_name/default.py --skip_train & # export CUDA_VISIBLE_DEVICES=3&&python render.py --model_path output/$exp_name/hand1-dense-v2 --configs arguments/$exp_name/hand1-dense-v2.py --skip_train wait -export CUDA_VISIBLE_DEVICES=0&&python metrics.py --model_path "output/$exp_name/cut_roasted_beefe/" & +export CUDA_VISIBLE_DEVICES=0&&python render.py --model_path output/$exp_name/flame_salmon_1 --configs arguments/$exp_name/default.py --skip_train & +export CUDA_VISIBLE_DEVICES=1&&python render.py --model_path output/$exp_name/flame_steak --configs arguments/$exp_name/default.py --skip_train & +export CUDA_VISIBLE_DEVICES=2&&python render.py --model_path output/$exp_name/coffee_martini --configs arguments/$exp_name/default.py --skip_train & +wait +export CUDA_VISIBLE_DEVICES=0&&python metrics.py --model_path "output/$exp_name/cut_roasted_beef/" & export CUDA_VISIBLE_DEVICES=1&&python metrics.py --model_path "output/$exp_name/cook_spinach/" & export CUDA_VISIBLE_DEVICES=2&&python metrics.py --model_path "output/$exp_name/sear_steak/" & # export CUDA_VISIBLE_DEVICES=3&&python metrics.py --model_path "output/$exp_name/hand1-dense-v2/" wait +export CUDA_VISIBLE_DEVICES=0&&python metrics.py --model_path "output/$exp_name/flame_salmon_1/" & +export CUDA_VISIBLE_DEVICES=1&&python metrics.py --model_path "output/$exp_name/flame_steak/" & +export CUDA_VISIBLE_DEVICES=2&&python metrics.py --model_path "output/$exp_name/coffee_martini/" & echo "Done" \ No newline at end of file diff --git a/train.py b/train.py index 20f931e..9f6b0a1 100644 --- a/train.py +++ b/train.py @@ -89,7 +89,7 @@ def scene_reconstruction(dataset, opt, hyper, pipe, testing_iterations, saving_i # Pick a random Camera if not viewpoint_stack: viewpoint_stack = scene.getTrainCameras() - batch_size = 1 + batch_size = opt.batch_size viewpoint_stack_loader = DataLoader(viewpoint_stack, batch_size=batch_size,shuffle=True,num_workers=32,collate_fn=list) loader = iter(viewpoint_stack_loader) if opt.dataloader: @@ -97,7 +97,7 @@ def scene_reconstruction(dataset, opt, hyper, pipe, testing_iterations, saving_i viewpoint_cams = next(loader) except StopIteration: print("reset dataloader") - batch_size = 1 + batch_size = opt.batch_size loader = iter(viewpoint_stack_loader) else: idx = randint(0, len(viewpoint_stack)-1)