423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-12-21 06:40:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

NeRF training: avoid caching unused visualization data.

Browse Source 
Summary: If we are not visualizing the training with visdom, then there are a couple of outputs of the coarse rendering step which are not small and are returned by the renderer but never used. We don't need to bother transferring them to the CPU.

Reviewed By: nikhilaravi

Differential Revision: D28939958

fbshipit-source-id: 7e0d6681d6524f7fb57b6b20164580006120de80
This commit is contained in:
Jeremy Reizenstein
2021-06-08 04:34:14 -07:00
committed by Facebook GitHub Bot
parent 7204a4ca64
commit f00ef66727
2 changed files with 23 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
projects/nerf/train_nerf.py
View File

@@ -52,6 +52,7 @@ def main(cfg: DictConfig):
n_hidden_neurons_dir=cfg.implicit_function.n_hidden_neurons_dir,
n_layers_xyz=cfg.implicit_function.n_layers_xyz,
density_noise_std=cfg.implicit_function.density_noise_std,
visualization=cfg.visualization.visdom,
)
# Move the model to the relevant device.
@@ -195,17 +196,18 @@ def main(cfg: DictConfig):
stats.print(stat_set="train")
# Update the visualization cache.
visuals_cache.append(
{
"camera": camera.cpu(),
"camera_idx": camera_idx,
"image": image.cpu().detach(),
"rgb_fine": nerf_out["rgb_fine"].cpu().detach(),
"rgb_coarse": nerf_out["rgb_coarse"].cpu().detach(),
"rgb_gt": nerf_out["rgb_gt"].cpu().detach(),
"coarse_ray_bundle": nerf_out["coarse_ray_bundle"],
}
)
if viz is not None:
visuals_cache.append(
{
"camera": camera.cpu(),
"camera_idx": camera_idx,
"image": image.cpu().detach(),
"rgb_fine": nerf_out["rgb_fine"].cpu().detach(),
"rgb_coarse": nerf_out["rgb_coarse"].cpu().detach(),
"rgb_gt": nerf_out["rgb_gt"].cpu().detach(),
"coarse_ray_bundle": nerf_out["coarse_ray_bundle"],
}
)
# Adjust the learning rate.
lr_scheduler.step()