feat: Update sampler for train dataloader.

Author: RoyYang0714

vis4d/data/loader.py

@@ -8,7 +8,12 @@
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import (
+    DataLoader,
+    Dataset,
+    RandomSampler,
+    SequentialSampler,
+)
 from torch.utils.data.distributed import DistributedSampler, Sampler
 
 from vis4d.common.distributed import get_rank, get_world_size
@@ -167,20 +172,29 @@ def _worker_init_fn(worker_id: int) -> None:
             if disable_subprocess_warning and worker_id != 0:
                 warnings.simplefilter("ignore")
 
-    if get_world_size() > 1 and sampler is None:
-        sampler = DistributedSampler(
-            dataset, shuffle=shuffle, drop_last=drop_last
-        )
-        shuffle = False
-        drop_last = False
+    if sampler is None:
+        if get_world_size() > 1:
+            sampler = DistributedSampler(
+                dataset, shuffle=shuffle, drop_last=drop_last
+            )
+            shuffle = False
+            drop_last = False
+        else:
+            if shuffle:
+                sampler = RandomSampler(dataset)
+                shuffle = False
+            else:
+                sampler = SequentialSampler(dataset)
 
     batch_sampler = None
-    if aspect_ratio_grouping and sampler is not None:
+    if aspect_ratio_grouping:
         batch_sampler = AspectRatioBatchSampler(
-            sampler, batch_size=samples_per_gpu
+            sampler, batch_size=samples_per_gpu, drop_last=drop_last
         )
         samples_per_gpu = 1
         shuffle = None
+        drop_last = False
+        sampler = None
 
     dataloader = DataLoader(
         dataset,
@@ -189,7 +203,7 @@ def _worker_init_fn(worker_id: int) -> None:
         collate_fn=(
             _collate_fn_multi if dataset.has_reference else _collate_fn_single
         ),
-        sampler=sampler if not aspect_ratio_grouping else None,
+        sampler=sampler,
         batch_sampler=batch_sampler,
         worker_init_fn=_worker_init_fn,
         persistent_workers=workers_per_gpu > 0,

vis4d/data/samplers.py

@@ -115,7 +115,14 @@ def __init__(
 
     def __iter__(self):
         for idx in self.sampler:
-            data_dict = self.sampler.dataset[idx]
+            if hasattr(self.sampler, "dataset"):
+                data_dict = self.sampler.dataset[idx]
+            elif hasattr(self.sampler, "data_source"):
+                data_dict = self.sampler.data_source[idx]
+            else:
+                raise ValueError(
+                    "sampler should have dataset or data_source attribute"
+                )
             height, width = data_dict[K.input_hw]
             bucket_id = 0 if width < height else 1
             bucket = self._aspect_ratio_buckets[bucket_id]

vis4d/engine/callbacks/visualizer.py

@@ -137,7 +137,6 @@ def on_test_batch_end(
                     trainer.logger.log_image(
                         key=f"{self.visualizer}/{cur_iter}",
                         images=[image],
-                        step=trainer.global_step,
                     )
             synchronize()