Updates on examples (#174)

* move torch data transfer into dataloader

* Update README

* use args.data_root

* Remove redundant check

* Fix isort

* Fix black

Author: 97littleleaf11

README.md

@@ -132,7 +132,7 @@ optimizer.step()
 
 ## Examples: 
 
-Before running those example scripts, please check the script about which dataset it is needed, and download the dataset first.
+Before running those example scripts, please check the script about which dataset is needed, and download the dataset first. You could use `--data_root` to specify the path.
 
 ```bash
 # clone the repo with submodules.

examples/datasets/dnerf_synthetic.py

@@ -86,6 +86,7 @@ def __init__(
         near: float = None,
         far: float = None,
         batch_over_images: bool = True,
+        device: str = "cuda:0",
     ):
         super().__init__()
         assert split in self.SPLITS, "%s" % split
@@ -106,18 +107,23 @@ def __init__(
             self.focal,
             self.timestamps,
         ) = _load_renderings(root_fp, subject_id, split)
-        self.images = torch.from_numpy(self.images).to(torch.uint8)
-        self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
-        self.timestamps = torch.from_numpy(self.timestamps).to(torch.float32)[
-            :, None
-        ]
+        self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
+        self.camtoworlds = (
+            torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
+        )
+        self.timestamps = (
+            torch.from_numpy(self.timestamps)
+            .to(device)
+            .to(torch.float32)[:, None]
+        )
         self.K = torch.tensor(
             [
                 [self.focal, 0, self.WIDTH / 2.0],
                 [0, self.focal, self.HEIGHT / 2.0],
                 [0, 0, 1],
             ],
             dtype=torch.float32,
+            device=device,
         )  # (3, 3)
         assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH)
 

examples/datasets/nerf_360_v2.py

@@ -169,6 +169,7 @@ def __init__(
         far: float = None,
         batch_over_images: bool = True,
         factor: int = 1,
+        device: str = "cuda:0",
     ):
         super().__init__()
         assert split in self.SPLITS, "%s" % split
@@ -186,9 +187,11 @@ def __init__(
         self.images, self.camtoworlds, self.K = _load_colmap(
             root_fp, subject_id, split, factor
         )
-        self.images = torch.from_numpy(self.images).to(torch.uint8)
-        self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
-        self.K = torch.tensor(self.K).to(torch.float32)
+        self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
+        self.camtoworlds = (
+            torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
+        )
+        self.K = torch.tensor(self.K).to(device).to(torch.float32)
         self.height, self.width = self.images.shape[1:3]
 
     def __len__(self):

examples/datasets/nerf_synthetic.py

@@ -79,6 +79,7 @@ def __init__(
         near: float = None,
         far: float = None,
         batch_over_images: bool = True,
+        device: str = "cuda:0",
     ):
         super().__init__()
         assert split in self.SPLITS, "%s" % split
@@ -109,15 +110,18 @@ def __init__(
             self.images, self.camtoworlds, self.focal = _load_renderings(
                 root_fp, subject_id, split
             )
-        self.images = torch.from_numpy(self.images).to(torch.uint8)
-        self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
+        self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
+        self.camtoworlds = (
+            torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
+        )
         self.K = torch.tensor(
             [
                 [self.focal, 0, self.WIDTH / 2.0],
                 [0, self.focal, self.HEIGHT / 2.0],
                 [0, 0, 1],
             ],
             dtype=torch.float32,
+            device=device,
         )  # (3, 3)
         assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH)
 

examples/train_mlp_dnerf.py

@@ -4,7 +4,7 @@
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -24,6 +24,12 @@
     set_random_seed(42)
 
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_root",
+        type=str,
+        default=str(pathlib.Path.cwd() / "data/dnerf"),
+        help="the root dir of the dataset",
+    )
     parser.add_argument(
         "--train_split",
         type=str,
@@ -91,31 +97,22 @@
         gamma=0.33,
     )
     # setup the dataset
-    data_root_fp = "/home/ruilongli/data/dnerf/"
     target_sample_batch_size = 1 << 16
     grid_resolution = 128
 
     train_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split=args.train_split,
         num_rays=target_sample_batch_size // render_n_samples,
     )
-    train_dataset.images = train_dataset.images.to(device)
-    train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
-    train_dataset.K = train_dataset.K.to(device)
-    train_dataset.timestamps = train_dataset.timestamps.to(device)
 
     test_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split="test",
         num_rays=None,
     )
-    test_dataset.images = test_dataset.images.to(device)
-    test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
-    test_dataset.K = test_dataset.K.to(device)
-    test_dataset.timestamps = test_dataset.timestamps.to(device)
 
     occupancy_grid = OccupancyGrid(
         roi_aabb=args.aabb,
@@ -191,7 +188,7 @@
                     f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
                 )
 
-            if step >= 0 and step % max_steps == 0 and step > 0:
+            if step > 0 and step % max_steps == 0:
                 # evaluation
                 radiance_field.eval()
 

examples/train_mlp_nerf.py

@@ -4,7 +4,7 @@
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -23,6 +23,12 @@
     set_random_seed(42)
 
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_root",
+        type=str,
+        default=str(pathlib.Path.cwd() / "data/nerf_synthetic"),
+        help="the root dir of the dataset",
+    )
     parser.add_argument(
         "--train_split",
         type=str,
@@ -112,40 +118,31 @@
     if args.scene == "garden":
         from datasets.nerf_360_v2 import SubjectLoader
 
-        data_root_fp = "/home/ruilongli/data/360_v2/"
         target_sample_batch_size = 1 << 16
         train_dataset_kwargs = {"color_bkgd_aug": "random", "factor": 4}
         test_dataset_kwargs = {"factor": 4}
         grid_resolution = 128
     else:
         from datasets.nerf_synthetic import SubjectLoader
 
-        data_root_fp = "/home/ruilongli/data/nerf_synthetic/"
         target_sample_batch_size = 1 << 16
         grid_resolution = 128
 
     train_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split=args.train_split,
         num_rays=target_sample_batch_size // render_n_samples,
         **train_dataset_kwargs,
     )
 
-    train_dataset.images = train_dataset.images.to(device)
-    train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
-    train_dataset.K = train_dataset.K.to(device)
-
     test_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split="test",
         num_rays=None,
         **test_dataset_kwargs,
     )
-    test_dataset.images = test_dataset.images.to(device)
-    test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
-    test_dataset.K = test_dataset.K.to(device)
 
     occupancy_grid = OccupancyGrid(
         roi_aabb=args.aabb,
@@ -217,7 +214,7 @@
                     f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
                 )
 
-            if step >= 0 and step % max_steps == 0 and step > 0:
+            if step > 0 and step % max_steps == 0:
                 # evaluation
                 radiance_field.eval()
 

examples/train_ngp_nerf.py

@@ -4,7 +4,7 @@
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -23,6 +23,12 @@
     set_random_seed(42)
 
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_root",
+        type=str,
+        default=str(pathlib.Path.cwd() / "data"),
+        help="the root dir of the dataset",
+    )
     parser.add_argument(
         "--train_split",
         type=str,
@@ -87,40 +93,31 @@
     if args.unbounded:
         from datasets.nerf_360_v2 import SubjectLoader
 
-        data_root_fp = "/home/ruilongli/data/360_v2/"
         target_sample_batch_size = 1 << 20
         train_dataset_kwargs = {"color_bkgd_aug": "random", "factor": 4}
         test_dataset_kwargs = {"factor": 4}
         grid_resolution = 256
     else:
         from datasets.nerf_synthetic import SubjectLoader
 
-        data_root_fp = "/home/ruilongli/data/nerf_synthetic/"
         target_sample_batch_size = 1 << 18
         grid_resolution = 128
 
     train_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split=args.train_split,
         num_rays=target_sample_batch_size // render_n_samples,
         **train_dataset_kwargs,
     )
 
-    train_dataset.images = train_dataset.images.to(device)
-    train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
-    train_dataset.K = train_dataset.K.to(device)
-
     test_dataset = SubjectLoader(
         subject_id=args.scene,
-        root_fp=data_root_fp,
+        root_fp=args.data_root,
         split="test",
         num_rays=None,
         **test_dataset_kwargs,
     )
-    test_dataset.images = test_dataset.images.to(device)
-    test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
-    test_dataset.K = test_dataset.K.to(device)
 
     if args.auto_aabb:
         camera_locs = torch.cat(
@@ -260,7 +257,7 @@ def occ_eval_fn(x):
                     f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
                 )
 
-            if step >= 0 and step % max_steps == 0 and step > 0:
+            if step > 0 and step % max_steps == 0:
                 # evaluation
                 radiance_field.eval()