remove ddp when using a single gpu

Author: jeongyw12382

dataloader/blender.py

@@ -57,12 +57,12 @@ def train_dataloader(self):
             import torch_xla.core.xla_model as xm
             sampler = SingleImageDDPSampler(
                 self.batch_size, xm.xrt_world_size(), xm.get_ordinal(),
-                len(self.i_train), self.image_len, self.args.i_validation
+                len(self.i_train), self.image_len, self.args.i_validation, True
             )
         else:
             sampler = SingleImageDDPSampler(
                 self.batch_size, None, None, len(self.i_train), 
-                self.image_len, self.args.i_validation
+                self.image_len, self.args.i_validation, False
             )
         return DataLoader(
             self.train_dset, batch_sampler=sampler, num_workers=self.args.num_workers,
@@ -73,9 +73,13 @@ def val_dataloader(self):
 
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.val_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.val_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.val_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.val_dset), False
+            )
 
         return DataLoader(
             self.val_dset, batch_size=self.chunk, sampler=sampler, 
@@ -86,9 +90,13 @@ def test_dataloader(self):
 
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.test_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.test_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.test_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.test_dset), False
+            )
 
         return DataLoader(
             self.test_dset, batch_size=self.chunk, sampler=sampler,
@@ -98,9 +106,13 @@ def test_dataloader(self):
     def predict_dataloader(self):
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.predict_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.predict_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.predict_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.predict_dset), False
+            )
 
         return DataLoader(
             self.predict_dset, batch_size=self.args.chunk, sampler=sampler,

dataloader/llff.py

@@ -61,11 +61,11 @@ def train_dataloader(self):
             import torch_xla.core.xla_model as xm
             sampler = MultipleImageDDPSampler(
                 self.batch_size, xm.xrt_world_size(), xm.get_ordinal(),
-                len(self.train_dset), self.args.i_validation
+                len(self.train_dset), self.args.i_validation, True
             )
         else:
             sampler = MultipleImageDDPSampler(
-                self.batch_size, None, None, len(self.train_dset), self.args.i_validation
+                self.batch_size, None, None, len(self.train_dset), self.args.i_validation, False
             )
 
         return DataLoader(
@@ -77,9 +77,13 @@ def val_dataloader(self):
 
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.val_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.val_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.val_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.val_dset), False
+            )
 
         return DataLoader(
             self.val_dset, batch_size=self.chunk, sampler=sampler, 
@@ -90,9 +94,13 @@ def test_dataloader(self):
 
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.test_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.test_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.test_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.test_dset), False
+            )
 
         return DataLoader(
             self.test_dset, batch_size=self.chunk, sampler=sampler, 
@@ -102,9 +110,13 @@ def test_dataloader(self):
     def predict_dataloader(self):
         if self.args.tpu: 
             import torch_xla.core.xla_model as xm
-            sampler = DDPSequnetialSampler(self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.predict_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, xm.xrt_world_size(), xm.get_ordinal(), len(self.predict_dset), True
+            )
         else:
-            sampler = DDPSequnetialSampler(self.chunk, None, None, len(self.predict_dset))
+            sampler = DDPSequnetialSampler(
+                self.chunk, None, None, len(self.predict_dset), False
+            )
 
         return DataLoader(
             self.predict_dset, batch_size=self.chunk, sampler=sampler,

dataloader/sampler.py

@@ -7,19 +7,23 @@
 
 
 class DDPSequnetialSampler(SequentialSampler): 
-    def __init__(self, batch_size, num_replicas, rank, N_total):
+    def __init__(self, batch_size, num_replicas, rank, N_total, tpu):
         self.data_source=None
         self.batch_size = batch_size
         self.N_total = N_total
         self.drop_last=False
-        if num_replicas is None:
-            if not dist.is_available():
-                raise RuntimeError("Requires distributed package to be available")
-            num_replicas = dist.get_world_size()
-        if rank is None:
-            if not dist.is_available():
-                raise RuntimeError("Requires distributed package to be available")
-            rank = dist.get_rank()
+        ngpus = torch.cuda.device_count()
+        if ngpus == 1 and not tpu:
+            rank, num_replicas = 0, 1
+        else:
+            if num_replicas is None:
+                if not dist.is_available():
+                    raise RuntimeError("Requires distributed package to be available")
+                num_replicas = dist.get_world_size()
+            if rank is None:
+                if not dist.is_available():
+                    raise RuntimeError("Requires distributed package to be available")
+                rank = dist.get_rank()
         self.rank = rank
         self.num_replicas = num_replicas
 
@@ -32,20 +36,24 @@ def __len__(self):
 
 
 class SingleImageDDPSampler:
-    def __init__(self, batch_size, num_replicas, rank, N_img, N_pixels, i_validation):
+    def __init__(self, batch_size, num_replicas, rank, N_img, N_pixels, i_validation, tpu):
         self.batch_size = batch_size
         self.N_pixels = N_pixels
         self.N_img = N_img
         self.drop_last = False
         self.i_validation = i_validation
-        if num_replicas is None:
-            if not dist.is_available():
-                raise RuntimeError("Requires distributed package to be available")
-            num_replicas = dist.get_world_size()
-        if rank is None:
-            if not dist.is_available():
-                raise RuntimeError("Require distributed package to be available")
-            rank = dist.get_rank()
+        ngpus = torch.cuda.device_count()
+        if ngpus == 1 and not tpu:
+            rank, num_replicas = 0, 1
+        else:
+            if num_replicas is None:
+                if not dist.is_available():
+                    raise RuntimeError("Requires distributed package to be available")
+                num_replicas = dist.get_world_size()
+            if rank is None:
+                if not dist.is_available():
+                    raise RuntimeError("Require distributed package to be available")
+                rank = dist.get_rank()
         self.rank = rank
         self.num_replicas = num_replicas
 
@@ -66,21 +74,25 @@ def __len__(self):
 
 
 class MultipleImageDDPSampler(DistributedSampler):
-    def __init__(self, batch_size, num_replicas, rank, total_len, i_validation):
+    def __init__(self, batch_size, num_replicas, rank, total_len, i_validation, tpu):
         self.batch_size = batch_size
         self.total_len = total_len
         self.i_validation = i_validation
-        self.drop_last = False
-        if num_replicas is None:
-            if not dist.is_available():
-                raise RuntimeError(
-                    "Require distributed package to be available")
-            num_replicas = dist.get_world_size()
-        if rank is None:
-            if not dist.is_available():
-                raise RuntimeError(
-                    "Require distributed package to be available")
-            rank = dist.get_rank()
+        self.drop_last = False        
+        ngpus = torch.cuda.device_count()
+        if ngpus == 1 and not tpu:
+            rank, num_replicas = 0, 1
+        else:
+            if num_replicas is None:
+                if not dist.is_available():
+                    raise RuntimeError(
+                        "Require distributed package to be available")
+                num_replicas = dist.get_world_size()
+            if rank is None:
+                if not dist.is_available():
+                    raise RuntimeError(
+                        "Require distributed package to be available")
+                rank = dist.get_rank()
         self.num_replicas = num_replicas
         self.rank = rank
 

run.py

@@ -69,8 +69,8 @@
         tpu_cores=args.tpu_num if args.tpu else None,
         replace_sampler_ddp=False,
         deterministic=True,
-        strategy=DDPPlugin(
-            find_unused_parameters=False) if not args.tpu else None,
+        strategy=DDPPlugin(find_unused_parameters=False) \
+            if n_gpus > 1 and not args.tpu else None,
         check_val_every_n_epoch=1,
         precision=32,
         num_sanity_val_steps=0,

scripts/jaxnerf_torch/eval.sh

@@ -1,72 +1,72 @@
 
 # LLFF
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/fern \
-    --expname fern --eval --model jaxnerf_torch
+    --expname fern --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/flower \
-    --expname flower --eval --model jaxnerf_torch
+    --expname flower --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/fortress \
-    --expname fortress --eval --model jaxnerf_torch
+    --expname fortress --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/horns \
-    --expname horns --eval --model jaxnerf_torch
+    --expname horns --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/leaves \
-    --expname leaves --eval --model jaxnerf_torch
+    --expname leaves --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/orchids \
-    --expname orchids --eval --model jaxnerf_torch
+    --expname orchids --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/room \
-    --expname room --eval --model jaxnerf_torch
+    --expname room --eval 
 python3 -m run --config configs/jaxnerf_torch/llff.yaml --datadir data/llff/trex \
-    --expname trex --eval --model jaxnerf_torch
+    --expname trex --eval 
 
 # LLFF_LARGE
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/fern \
-    --expname fern_large --eval --model jaxnerf_torch
+    --expname fern_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/flower \
-    --expname flower_large --eval --model jaxnerf_torch
+    --expname flower_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/fortress \
-    --expname fortress_large --eval --model jaxnerf_torch
+    --expname fortress_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/horns \
-    --expname horns_large --eval --model jaxnerf_torch
+    --expname horns_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/leaves \
-    --expname leaves_large --eval --model jaxnerf_torch
+    --expname leaves_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/orchids \
-    --expname orchids_large --eval --model jaxnerf_torch
+    --expname orchids_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/room \
-    --expname room_large --eval --model jaxnerf_torch
+    --expname room_large --eval 
 python3 -m run --config configs/jaxnerf_torch/llff_large.yaml --datadir data/llff/trex \
-    --expname trex_large --eval --model jaxnerf_torch
+    --expname trex_large --eval 
 
 # BLENDER
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/chair \
-    --expname chair --eval --model jaxnerf_torch
+    --expname chair --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/drums \
-    --expname drums --eval --model jaxnerf_torch
+    --expname drums --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/ficus \
-    --expname ficus --eval --model jaxnerf_torch
+    --expname ficus --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/hotdog \
-    --expname hotdog --eval --model jaxnerf_torch
+    --expname hotdog --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/lego \
-    --expname lego --eval --model jaxnerf_torch
+    --expname lego --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/materials \
-    --expname materials --eval --model jaxnerf_torch
+    --expname materials --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/mic \
-    --expname mic --eval --model jaxnerf_torch
+    --expname mic --eval 
 python3 -m run --config configs/jaxnerf_torch/blender.yaml --datadir data/blender/ship \
-    --expname ship --eval --model jaxnerf_torch
+    --expname ship --eval 
 
 # BLENDER LARGE
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/chair \
-    --expname chair_large --eval --model jaxnerf_torch
+    --expname chair_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/drums \
-    --expname drums_large --eval --model jaxnerf_torch
+    --expname drums_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/ficus \
-    --expname ficus_large --eval --model jaxnerf_torch
+    --expname ficus_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/hotdog \
-    --expname hotdog_large --eval --model jaxnerf_torch
+    --expname hotdog_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/lego \
-    --expname lego_large --eval --model jaxnerf_torch
+    --expname lego_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/materials \
-    --expname materials_large --eval --model jaxnerf_torch
+    --expname materials_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/mic \
-    --expname mic_large --eval --model jaxnerf_torch
+    --expname mic_large --eval 
 python3 -m run --config configs/jaxnerf_torch/blender_large.yaml --datadir data/blender/ship \
-    --expname ship_large --eval --model jaxnerf_torch
+    --expname ship_large --eval