diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 979f8f986..f0eeb80e2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -83,7 +83,7 @@ method_configs["nerfacto"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=16000, + train_num_rays_per_batch=17000, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 388de7fb4..e9dd0f3fe 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -351,7 +351,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method # Otherwise, use the default pixel sampler if is_equirectangular.any(): CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - return PatchPixelSampler(*args, **kwargs,patch_size=8) + return PatchPixelSampler(*args, **kwargs,patch_size=16) def setup_train(self): """Sets up the data loaders for training""" diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index f64c18a9f..228f4bb34 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -303,26 +303,28 @@ def depth_loss( raise NotImplementedError("Provided depth loss type not implemented.") -def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=8): +def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=16): out_patches = rgb_outputs.view(-1,1,patch_size,patch_size,3) gt_patches = rgb_gt.view(-1,1,patch_size,patch_size,3) device = rgb_outputs.device batch_size = out_patches.shape[0] residuals = torch.mean((out_patches - gt_patches)**2,dim=-1) - # with torch.no_grad(): - med_residual = torch.median(residuals) - #equation 8 - weight = (residuals<=med_residual).float() - #equation 9 - blurred_w = (torch.nn.functional.conv2d(weight,(1/9.)*torch.ones((1,1,3,3),device=device),padding='same')>=0.5).float() - # expected_w = torch.mean(torch.mean(blurred_w,dim=-1),dim=-1) - #equation 10 - # weight_r8 = (expected_w > 0.6).float() - # w_8x8 = weight_r8[...,None] * torch.ones((batch_size,8,8),device=device) - # final_w = torch.zeros((batch_size,patch_size,patch_size),device=device) - # final_w[:, - # (patch_size//2-patch_size//4):(patch_size//2+patch_size//4), - # (patch_size//2-patch_size//4):(patch_size//2+patch_size//4)] = w_8x8 - # loss = torch.mean(residuals.squeeze()*blurred_w) - loss = torch.mean(residuals.squeeze()*blurred_w) + with torch.no_grad(): + med_residual = torch.quantile(residuals,.9) + #equation 8 + weight = (residuals<=med_residual).float()#B x patch x patch + #equation 9 + weight = torch.nn.functional.pad(weight,(1,1,1,1),mode='replicate') + blurred_w = (torch.nn.functional.conv2d(weight,(1/9.)*torch.ones((1,1,3,3),device=device), + padding='valid')>=0.5).float() + #equation 10 + expected_w = blurred_w.view(batch_size,-1).mean(1) + weight_r8 = (expected_w >= 0.6).float() #Bx1 + #the paper uses the lines below and multiplies residuals by final_w, but I found that just using + #the value for all 16 pixels worked better. + # final_w = torch.zeros((batch_size,patch_size,patch_size),device=device) + # final_w[:, + # (patch_size//2-patch_size//4):(patch_size//2+patch_size//4), + # (patch_size//2-patch_size//4):(patch_size//2+patch_size//4)] = weight_r8[:,None,None] + # loss = torch.mean(residuals.squeeze()*weight_r8[:,None,None]) return loss \ No newline at end of file