Add support mip nerf in training example

[foreverYoungGitHub]

Hi, @liruilong940607,

Thanks for your awesome work!

I just add the support for training the mip nerf (kind of) with nerfacc based on your discussion at #88. I did a quick test for the lego scene with 1 t4 and this is the full logs:

elapsed_time=1.76s | step=0 | loss=0.07456 | alive_ray_mask=64 | n_rendering_samples=10341 | num_rays=64 |
elapsed_time=675.89s | step=5000 | loss=0.00614 | alive_ray_mask=1972 | n_rendering_samples=65941 | num_rays=6340 |
elapsed_time=1436.05s | step=10000 | loss=0.00325 | alive_ray_mask=2624 | n_rendering_samples=62479 | num_rays=8388 |
elapsed_time=2252.93s | step=15000 | loss=0.00259 | alive_ray_mask=3201 | n_rendering_samples=65802 | num_rays=10022 |
elapsed_time=3095.50s | step=20000 | loss=0.00191 | alive_ray_mask=3379 | n_rendering_samples=65854 | num_rays=10483 |
elapsed_time=3944.66s | step=25000 | loss=0.00169 | alive_ray_mask=3569 | n_rendering_samples=66453 | num_rays=11441 |
elapsed_time=4818.70s | step=30000 | loss=0.00130 | alive_ray_mask=3888 | n_rendering_samples=65518 | num_rays=12459 |
elapsed_time=5703.86s | step=35000 | loss=0.00117 | alive_ray_mask=4004 | n_rendering_samples=65890 | num_rays=12501 |
elapsed_time=6595.80s | step=40000 | loss=0.00116 | alive_ray_mask=4163 | n_rendering_samples=66805 | num_rays=13105 |
elapsed_time=7495.55s | step=45000 | loss=0.00095 | alive_ray_mask=4178 | n_rendering_samples=66799 | num_rays=13154 |
elapsed_time=8396.90s | step=50000 | loss=0.00086 | alive_ray_mask=4242 | n_rendering_samples=67710 | num_rays=13257 |
100%|██████████| 200/200 [19:52<00:00,  5.96s/it]
evaluation: psnr_avg=33.62006133079529

The training time is pretty similar with vanilla nerf and the performance is almost same as well (which I do expect slightly better performance).

It would be nice if you can take a look at it when you are available!

Thank you in advance!

Btw, this pr also changes the dataset root dir to an argument for easier usage 😉

[foreverYoungGitHub]

@liruilong940607 @tancik I'm wondering if you can take a look at it when you are available?

[liruilong940607]

Hey thanks for the implementation! I quickly skimmed it and it looks all reasonable! I will take a close look in the weekend.

[foreverYoungGitHub]

That would be super helpful! Thank you Li!