Updates on examples (#174)

* move torch data transfer into dataloader

* Update README

* use args.data_root

* Remove redundant check

* Fix isort

* Fix black

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 @@ class SubjectLoader(torch.utils.data.Dataset):
         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,11 +107,15 @@ class SubjectLoader(torch.utils.data.Dataset):
             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],
@@ -118,6 +123,7 @@ class SubjectLoader(torch.utils.data.Dataset):
                 [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 @@ class SubjectLoader(torch.utils.data.Dataset):
         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 @@ class SubjectLoader(torch.utils.data.Dataset):
         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 @@ class SubjectLoader(torch.utils.data.Dataset):
         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,8 +110,10 @@ class SubjectLoader(torch.utils.data.Dataset):
             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],
@@ -118,6 +121,7 @@ class SubjectLoader(torch.utils.data.Dataset):
                 [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 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -24,6 +24,12 @@ if __name__ == "__main__":
     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 @@ if __name__ == "__main__":
         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 @@ if __name__ == "__main__":
                     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 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -23,6 +23,12 @@ if __name__ == "__main__":
     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,7 +118,6 @@ if __name__ == "__main__":
     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}
@@ -120,32 +125,24 @@ if __name__ == "__main__":
     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 @@ if __name__ == "__main__":
                     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 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
 
 import argparse
 import math
-import os
+import pathlib
 import time
 
 import imageio
@@ -23,6 +23,12 @@ if __name__ == "__main__":
     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,7 +93,6 @@ if __name__ == "__main__":
     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}
@@ -95,32 +100,24 @@ if __name__ == "__main__":
     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 @@ if __name__ == "__main__":
                     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()