tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Add dynamic batching to compound model (#535)

Browse Source
This commit is contained in:
Matthew Tancik 2022-09-24 21:40:23 -07:00 committed by GitHub
parent 06b1c80d2e
commit aea0cfc253
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfactory/configs/base_configs.py
View File

@ -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),
),

8
nerfactory/models/compound.py
View File

@ -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"]