diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index ab154dce4..b240af494 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -55,7 +55,7 @@ base_configs: Dict[str, Config] = {} base_configs["compound"] = Config( method_name="compound", trainer=TrainerConfig(mixed_precision=True), - pipeline=VanillaPipelineConfig( + pipeline=DynamicBatchPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), model=CompoundModelConfig(eval_num_rays_per_chunk=8192), ), diff --git a/nerfactory/models/compound.py b/nerfactory/models/compound.py index 3f9106ac4..398d361b6 100644 --- a/nerfactory/models/compound.py +++ b/nerfactory/models/compound.py @@ -186,9 +186,17 @@ class CompoundModel(Model): "accumulation": accumulation, "depth": depth, "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler + "num_samples_per_ray": packed_info[:, 1], } return outputs + def get_metrics_dict(self, outputs, batch): + image = batch["image"].to(self.device) + metrics_dict = {} + metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) + metrics_dict["num_samples_per_batch"] = outputs["num_samples_per_ray"].sum() + return metrics_dict + def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) mask = outputs["alive_ray_mask"]