vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/EDM/src/edm/utils/supervision.py

@@ -0,0 +1,255 @@
+from math import log
+from loguru import logger
+
+import torch
+from einops import repeat
+from kornia.utils import create_meshgrid
+
+from .geometry import warp_kpts
+
+##############  ↓  Coarse-Level supervision  ↓  ##############
+
+
+@torch.no_grad()
+def mask_pts_at_padded_regions(grid_pt, mask):
+    """For megadepth dataset, zero-padding exists in images"""
+    mask = repeat(mask, "n h w -> n (h w) c", c=2)
+    grid_pt[~mask.bool()] = 0
+    return grid_pt
+
+
+@torch.no_grad()
+def spvs_coarse(data, config):
+    """
+    Update:
+        data (dict): {
+            "conf_matrix_gt": [N, hw0, hw1],
+            'spv_b_ids': [M]
+            'spv_i_ids': [M]
+            'spv_j_ids': [M]
+            'spv_w_pt0_i': [N, hw0, 2], in original image resolution
+            'spv_pt1_i': [N, hw1, 2], in original image resolution
+        }
+
+    NOTE:
+        - for scannet dataset, there're 3 kinds of resolution {i, c, f}
+        - for megadepth dataset, there're 4 kinds of resolution {i, i_resize, c, f}
+    """
+    # 1. misc
+    device = data["image0"].device
+    N, _, H0, W0 = data["image0"].shape
+    _, _, H1, W1 = data["image1"].shape
+    scale = config["EDM"]["LOCAL_RESOLUTION"]
+
+    scale0 = scale * data["scale0"][:, None] if "scale0" in data else scale
+    scale1 = scale * data["scale1"][:, None] if "scale0" in data else scale
+    h0, w0, h1, w1 = map(lambda x: x // scale, [H0, W0, H1, W1])
+
+    # 2. warp grids
+    # create kpts in meshgrid and resize them to image resolution
+    grid_pt0_c = (
+        create_meshgrid(h0, w0, False, device).reshape(
+            1, h0 * w0, 2).repeat(N, 1, 1)
+    )  # [N, hw, 2]
+    grid_pt0_i = scale0 * grid_pt0_c
+    grid_pt1_c = (
+        create_meshgrid(h1, w1, False, device).reshape(
+            1, h1 * w1, 2).repeat(N, 1, 1)
+    )
+    grid_pt1_i = scale1 * grid_pt1_c
+
+    # mask padded region to (0, 0), so no need to manually mask conf_matrix_gt
+    if "mask0" in data:
+        grid_pt0_i = mask_pts_at_padded_regions(grid_pt0_i, data["mask0"])
+        grid_pt1_i = mask_pts_at_padded_regions(grid_pt1_i, data["mask1"])
+
+    # warp kpts bi-directionally and resize them to coarse-level resolution
+    # (no depth consistency check, since it leads to worse results experimentally)
+    # (unhandled edge case: points with 0-depth will be warped to the left-up corner)
+    _, w_pt0_i = warp_kpts(
+        grid_pt0_i,
+        data["depth0"],
+        data["depth1"],
+        data["T_0to1"],
+        data["K0"],
+        data["K1"],
+    )
+    _, w_pt1_i = warp_kpts(
+        grid_pt1_i,
+        data["depth1"],
+        data["depth0"],
+        data["T_1to0"],
+        data["K1"],
+        data["K0"],
+    )
+    w_pt0_c = w_pt0_i / scale1
+    w_pt1_c = w_pt1_i / scale0
+
+    # 3. check if mutual nearest neighbor
+    w_pt0_c_round = w_pt0_c[:, :, :].round().long()
+    nearest_index1 = w_pt0_c_round[..., 0] + w_pt0_c_round[..., 1] * w1
+    w_pt1_c_round = w_pt1_c[:, :, :].round().long()
+    nearest_index0 = w_pt1_c_round[..., 0] + w_pt1_c_round[..., 1] * w0
+
+    # corner case: out of boundary
+
+    def out_bound_mask(pt, w, h):
+        return (
+            (pt[..., 0] < 0) + (pt[..., 0] >= w) +
+            (pt[..., 1] < 0) + (pt[..., 1] >= h)
+        )
+
+    nearest_index1[out_bound_mask(w_pt0_c_round, w1, h1)] = 0
+    nearest_index0[out_bound_mask(w_pt1_c_round, w0, h0)] = 0
+
+    loop_back = torch.stack(
+        [nearest_index0[_b][_i] for _b, _i in enumerate(nearest_index1)], dim=0
+    )
+    correct_0to1 = loop_back == torch.arange(
+        h0 * w0, device=device)[None].repeat(N, 1)
+    correct_0to1[:, 0] = False  # ignore the top-left corner
+
+    # 4. construct a gt conf_matrix
+    conf_matrix_gt = torch.zeros(N, h0 * w0, h1 * w1, device=device)
+    b_ids, i_ids = torch.where(correct_0to1 != 0)
+    j_ids = nearest_index1[b_ids, i_ids]
+
+    conf_matrix_gt[b_ids, i_ids, j_ids] = 1
+    data.update({"conf_matrix_gt": conf_matrix_gt})
+
+    # 5. save coarse matches(gt) for training fine level
+    if len(b_ids) == 0:
+        logger.warning(
+            f"No groundtruth coarse match found for: {data['pair_names']}")
+        # this won't affect fine-level loss calculation
+        b_ids = torch.tensor([0], device=device)
+        i_ids = torch.tensor([0], device=device)
+        j_ids = torch.tensor([0], device=device)
+
+    data.update({"spv_b_ids": b_ids, "spv_i_ids": i_ids, "spv_j_ids": j_ids})
+
+    # 6. save intermediate results (for fast fine-level computation)
+    data.update(
+        {
+            "spv_w_pt0_i": w_pt0_i,
+            "spv_pt1_i": grid_pt1_i,
+            "spv_w_pt1_i": w_pt1_i,
+            "spv_pt0_i": grid_pt0_i,
+        }
+    )
+
+
+def compute_supervision_coarse(data, config):
+    assert (
+        len(set(data["dataset_name"])) == 1
+    ), "Do not support mixed datasets training!"
+    data_source = data["dataset_name"][0]
+    if data_source.lower() in ["scannet", "megadepth"]:
+        spvs_coarse(data, config)
+    else:
+        raise ValueError(f"Unknown data source: {data_source}")
+
+
+##############  ↓  Fine-Level supervision  ↓  ##############
+@torch.no_grad()
+def spvs_fine(data, pad=0, window_size=8.0):
+    """
+    Update:
+        data (dict):{
+            "expec_f_gt": [M, 2]}
+    """
+    # 1. get coarse prediction
+    w_pt0_i, pt1_i = data["spv_w_pt0_i"], data["spv_pt1_i"]
+    w_pt1_i, pt0_i = data["spv_w_pt1_i"], data["spv_pt0_i"]
+    b_ids, i_ids, j_ids, mconf = (
+        data["b_ids"].clone(),
+        data["i_ids"].clone(),
+        data["j_ids"].clone(),
+        data["mconf"].clone(),
+    )
+
+    scale0 = data["scale0"][b_ids] if "scale0" in data else 1.0
+    scale1 = data["scale1"][b_ids] if "scale0" in data else 1.0
+
+    # 2. random sampling of training samples for fine-level supervision
+    # (optional) In the early stages of training
+    # the coarse matching predictions are not good enough
+    # pad samples with gt coarse-level matches
+    if pad > 0:
+        bs = int(b_ids.max()) + 1
+        k = b_ids.shape[0] // bs
+        offset01_gt = (
+            w_pt0_i[b_ids, i_ids] - pt1_i[b_ids, j_ids]
+        ) / scale1  # [M, 2] [NK, 2]
+        offset01_gt_mask = (offset01_gt[:, 0].abs() < (window_size // 2)) & (
+            offset01_gt[:, 1].abs() < (window_size // 2)
+        )
+        spv_b_ids, spv_i_ids, spv_j_ids = (
+            data["spv_b_ids"].clone(),
+            data["spv_i_ids"].clone(),
+            data["spv_j_ids"].clone(),
+        )
+
+        # each pair in a batch
+        for bi in range(bs):
+            b_mask = b_ids == bi
+            spv_b_mask = spv_b_ids == bi
+            pad_num = 0
+            max_pad_num = pad
+            pred_true_num = offset01_gt_mask[b_mask].sum()
+
+            # Only pad when the number of coarse matches is insufficient
+            if pred_true_num < pad:
+                max_pad_num = pad - pred_true_num
+            else:
+                continue
+
+            candi_spv_i_ids = spv_i_ids[spv_b_mask]
+            candi_spv_j_ids = spv_j_ids[spv_b_mask]
+            shuffle_indices = torch.randperm(spv_b_mask.sum())
+            candi_spv_i_ids = candi_spv_i_ids[shuffle_indices]
+            candi_spv_j_ids = candi_spv_j_ids[shuffle_indices]
+
+            max_pad_num = min(int(spv_b_mask.sum()), max_pad_num)
+            for i in range(k):
+                if pad_num < max_pad_num:
+                    if offset01_gt_mask[b_mask].tolist()[i] is False:
+                        i_ids[bi * k + i] = candi_spv_i_ids[pad_num]
+                        j_ids[bi * k + i] = candi_spv_j_ids[pad_num]
+                        # Just assisting fine-level training
+                        mconf[bi * k + i] = 0.0
+                        pad_num += 1
+                else:
+                    break
+
+        # update data
+        data.update(
+            {
+                "b_ids": b_ids,
+                "i_ids": i_ids,
+                "j_ids": j_ids,
+                "mconf": mconf,
+            }
+        )
+
+    # 3. compute fine-level offset gt
+    offset01_gt = (
+        w_pt0_i[b_ids, i_ids] - pt1_i[b_ids, j_ids]
+    ) / scale1  # [M, 2] [NK, 2]
+    offset10_gt = (
+        w_pt1_i[b_ids, j_ids] - pt0_i[b_ids, i_ids]
+    ) / scale0  # [M, 2] [NK, 2]
+    data.update({"offset01_gt": offset01_gt})
+    data.update({"offset10_gt": offset10_gt})
+
+    target_uv = torch.cat([offset01_gt, offset10_gt])  # [2NK, 2] range [-4, 4]
+    pos_mask = (target_uv[:, 0].abs() < (window_size // 2)) & (
+        target_uv[:, 1].abs() < (window_size // 2)
+    )
+    # [M, 2] range [-0.5, 0.5]
+    data.update({"target_uv": target_uv / window_size})
+    data.update({"target_uv_weight": pos_mask})  # [M, 1]
+
+
+def compute_supervision_fine(data, pad=0, window_size=8.0):
+    spvs_fine(data, pad, window_size)

vismatch/third_party/EDM/src/lightning/data.py

@@ -0,0 +1,450 @@
+import os
+import math
+from collections import abc
+from loguru import logger
+from torch.utils.data.dataset import Dataset
+from tqdm import tqdm
+from os import path as osp
+from pathlib import Path
+from joblib import Parallel, delayed
+
+import lightning.pytorch as pl
+from torch import distributed as dist
+from torch.utils.data import (
+    Dataset,
+    DataLoader,
+    ConcatDataset,
+    DistributedSampler,
+    RandomSampler,
+    dataloader,
+)
+
+from src.utils.augment import build_augmentor
+from src.utils.dataloader import get_local_split
+from src.utils.misc import tqdm_joblib
+from src.utils import comm
+from src.datasets.megadepth import MegaDepthDataset
+from src.datasets.scannet import ScanNetDataset
+from src.datasets.sampler import RandomConcatSampler
+
+
+class MultiSceneDataModule(pl.LightningDataModule):
+    """
+    For distributed training, each training process is assgined
+    only a part of the training scenes to reduce memory overhead.
+    """
+
+    def __init__(self, args, config):
+        super().__init__()
+
+        # 1. data config
+        # Train and Val should from the same data source
+        self.trainval_data_source = config.DATASET.TRAINVAL_DATA_SOURCE
+        self.test_data_source = config.DATASET.TEST_DATA_SOURCE
+        # training and validating
+        self.train_data_root = config.DATASET.TRAIN_DATA_ROOT
+        self.train_pose_root = config.DATASET.TRAIN_POSE_ROOT  # (optional)
+        self.train_npz_root = config.DATASET.TRAIN_NPZ_ROOT
+        self.train_list_path = config.DATASET.TRAIN_LIST_PATH
+        self.train_intrinsic_path = config.DATASET.TRAIN_INTRINSIC_PATH
+        self.val_data_root = config.DATASET.VAL_DATA_ROOT
+        self.val_pose_root = config.DATASET.VAL_POSE_ROOT  # (optional)
+        self.val_npz_root = config.DATASET.VAL_NPZ_ROOT
+        self.val_list_path = config.DATASET.VAL_LIST_PATH
+        self.val_intrinsic_path = config.DATASET.VAL_INTRINSIC_PATH
+        # testing
+        self.test_data_root = config.DATASET.TEST_DATA_ROOT
+        self.test_pose_root = config.DATASET.TEST_POSE_ROOT  # (optional)
+        self.test_npz_root = config.DATASET.TEST_NPZ_ROOT
+        self.test_list_path = config.DATASET.TEST_LIST_PATH
+        self.test_intrinsic_path = config.DATASET.TEST_INTRINSIC_PATH
+
+        # 2. dataset config
+        # general options
+        self.min_overlap_score_test = (
+            config.DATASET.MIN_OVERLAP_SCORE_TEST
+        )  # 0.4, omit data with overlap_score < min_overlap_score
+        self.min_overlap_score_train = config.DATASET.MIN_OVERLAP_SCORE_TRAIN
+        self.augment_fn = build_augmentor(
+            config.DATASET.AUGMENTATION_TYPE
+        )  # None, options: [None, 'dark', 'mobile']
+
+        # ScanNet options
+
+        self.scan_img_resize = (config.EDM.TEST_RES_W,
+                                config.EDM.TEST_RES_H)  # (640, 480)
+
+        # MegaDepth options
+        self.mgdpt_img_resize = config.EDM.TEST_RES_W  # 832
+        self.mgdpt_img_pad = config.DATASET.MGDPT_IMG_PAD  # True
+        self.mgdpt_depth_pad = config.DATASET.MGDPT_DEPTH_PAD  # True
+        self.mgdpt_df = config.DATASET.MGDPT_DF  # 8
+        self.coarse_scale = 1 / config.EDM.LOCAL_RESOLUTION  # 1/8
+
+        self.fp16 = config.EDM.HALF
+
+        # 3.loader parameters
+        self.train_loader_params = {
+            "batch_size": args.batch_size,
+            "num_workers": args.num_workers,
+            "pin_memory": getattr(args, "pin_memory", True),
+        }
+        self.val_loader_params = {
+            "batch_size": 1,
+            "shuffle": False,
+            "num_workers": args.num_workers,
+            "pin_memory": getattr(args, "pin_memory", True),
+        }
+        self.test_loader_params = {
+            "batch_size": 1,
+            "shuffle": False,
+            "num_workers": args.num_workers,
+            "pin_memory": True,
+        }
+
+        # 4. sampler
+        self.data_sampler = config.TRAINER.DATA_SAMPLER
+        self.n_samples_per_subset = config.TRAINER.N_SAMPLES_PER_SUBSET
+        self.subset_replacement = config.TRAINER.SB_SUBSET_SAMPLE_REPLACEMENT
+        self.shuffle = config.TRAINER.SB_SUBSET_SHUFFLE
+        self.repeat = config.TRAINER.SB_REPEAT
+
+        # (optional) RandomSampler for debugging
+
+        # misc configurations
+        self.parallel_load_data = getattr(
+            args, "parallel_load_data", False)  # True for indoor
+        self.seed = config.TRAINER.SEED  # 66
+
+    def setup(self, stage=None):
+        """
+        Setup train / val / test dataset. This method will be called by PL automatically.
+        Args:
+            stage (str): 'fit' in training phase, and 'test' in testing phase.
+        """
+
+        assert stage in ["fit", "validate",
+                         "test"], "stage must be either fit or test"
+
+        try:
+            self.world_size = dist.get_world_size()
+            self.rank = dist.get_rank()
+            logger.info(f"[rank:{self.rank}] world_size: {self.world_size}")
+        except AssertionError as ae:
+            self.world_size = 1
+            self.rank = 0
+            # logger.warning(" (set wolrd_size=1 and rank=0)")
+            logger.warning(str(ae) + " (set wolrd_size=1 and rank=0)")
+
+        if stage == "fit":
+            self.train_dataset = self._setup_dataset(
+                self.train_data_root,
+                self.train_npz_root,
+                self.train_list_path,
+                self.train_intrinsic_path,
+                mode="train",
+                min_overlap_score=self.min_overlap_score_train,
+                pose_dir=self.train_pose_root,
+            )
+            # setup multiple (optional) validation subsets
+            if isinstance(self.val_list_path, (list, tuple)):
+                self.val_dataset = []
+                if not isinstance(self.val_npz_root, (list, tuple)):
+                    self.val_npz_root = [
+                        self.val_npz_root for _ in range(len(self.val_list_path))
+                    ]
+                for npz_list, npz_root in zip(self.val_list_path, self.val_npz_root):
+                    self.val_dataset.append(
+                        self._setup_dataset(
+                            self.val_data_root,
+                            npz_root,
+                            npz_list,
+                            self.val_intrinsic_path,
+                            mode="val",
+                            min_overlap_score=self.min_overlap_score_test,
+                            pose_dir=self.val_pose_root,
+                        )
+                    )
+            else:
+                self.val_dataset = self._setup_dataset(
+                    self.val_data_root,
+                    self.val_npz_root,
+                    self.val_list_path,
+                    self.val_intrinsic_path,
+                    mode="val",
+                    min_overlap_score=self.min_overlap_score_test,
+                    pose_dir=self.val_pose_root,
+                )
+            logger.info(f"[rank:{self.rank}] Train & Val Dataset loaded!")
+        elif stage == "validate":
+            if isinstance(self.val_list_path, (list, tuple)):
+                self.val_dataset = []
+                if not isinstance(self.val_npz_root, (list, tuple)):
+                    self.val_npz_root = [
+                        self.val_npz_root for _ in range(len(self.val_list_path))
+                    ]
+                for npz_list, npz_root in zip(self.val_list_path, self.val_npz_root):
+                    self.val_dataset.append(
+                        self._setup_dataset(
+                            self.val_data_root,
+                            npz_root,
+                            npz_list,
+                            self.val_intrinsic_path,
+                            mode="val",
+                            min_overlap_score=self.min_overlap_score_test,
+                            pose_dir=self.val_pose_root,
+                        )
+                    )
+            else:
+                self.val_dataset = self._setup_dataset(
+                    self.val_data_root,
+                    self.val_npz_root,
+                    self.val_list_path,
+                    self.val_intrinsic_path,
+                    mode="val",
+                    min_overlap_score=self.min_overlap_score_test,
+                    pose_dir=self.val_pose_root,
+                )
+            logger.info(f"[rank:{self.rank}] Val Dataset loaded!")
+        else:  # stage == 'test
+            self.test_dataset = self._setup_dataset(
+                self.test_data_root,
+                self.test_npz_root,
+                self.test_list_path,
+                self.test_intrinsic_path,
+                mode="test",
+                min_overlap_score=self.min_overlap_score_test,
+                pose_dir=self.test_pose_root,
+            )
+            logger.info(f"[rank:{self.rank}]: Test Dataset loaded!")
+
+    def _setup_dataset(
+        self,
+        data_root,
+        split_npz_root,
+        scene_list_path,
+        intri_path,
+        mode="train",
+        min_overlap_score=0.0,
+        pose_dir=None,
+    ):
+        """Setup train / val / test set"""
+        with open(scene_list_path, "r") as f:
+            npz_names = [name.split()[0] for name in f.readlines()]
+
+        if mode == "train":
+            local_npz_names = get_local_split(
+                npz_names, self.world_size, self.rank, self.seed
+            )
+        else:
+            local_npz_names = npz_names
+        logger.info(
+            f"[rank {self.rank}]: {len(local_npz_names)} scene(s) assigned.")
+
+        dataset_builder = (
+            self._build_concat_dataset_parallel
+            if self.parallel_load_data
+            else self._build_concat_dataset
+        )
+        return dataset_builder(
+            data_root,
+            local_npz_names,
+            split_npz_root,
+            intri_path,
+            mode=mode,
+            min_overlap_score=min_overlap_score,
+            pose_dir=pose_dir,
+        )
+
+    def _build_concat_dataset(
+        self,
+        data_root,
+        npz_names,
+        npz_dir,
+        intrinsic_path,
+        mode,
+        min_overlap_score=0.0,
+        pose_dir=None,
+    ):
+        datasets = []
+        augment_fn = self.augment_fn if mode == "train" else None
+        data_source = (
+            self.trainval_data_source
+            if mode in ["train", "val"]
+            else self.test_data_source
+        )
+        if str(data_source).lower() == "megadepth":
+            npz_names = [f"{n}.npz" for n in npz_names]
+        for npz_name in tqdm(
+            npz_names,
+            desc=f"[rank:{self.rank}] loading {mode} datasets",
+            disable=int(self.rank) != 0,
+        ):
+            # `ScanNetDataset`/`MegaDepthDataset` load all data from npz_path when initialized, which might take time.
+            npz_path = osp.join(npz_dir, npz_name)
+            if data_source == "ScanNet":
+                datasets.append(
+                    ScanNetDataset(
+                        data_root,
+                        npz_path,
+                        intrinsic_path,
+                        mode=mode,
+                        min_overlap_score=min_overlap_score,
+                        augment_fn=augment_fn,
+                        pose_dir=pose_dir,
+                        img_resize=self.scan_img_resize,
+                        fp16=self.fp16,
+                    )
+                )
+            elif data_source == "MegaDepth":
+                datasets.append(
+                    MegaDepthDataset(
+                        data_root,
+                        npz_path,
+                        mode=mode,
+                        min_overlap_score=min_overlap_score,
+                        img_resize=self.mgdpt_img_resize,
+                        df=self.mgdpt_df,
+                        img_padding=self.mgdpt_img_pad,
+                        depth_padding=self.mgdpt_depth_pad,
+                        augment_fn=augment_fn,
+                        coarse_scale=self.coarse_scale,
+                        fp16=self.fp16,
+                    )
+                )
+            else:
+                raise NotImplementedError()
+        return ConcatDataset(datasets)
+
+    def _build_concat_dataset_parallel(
+        self,
+        data_root,
+        npz_names,
+        npz_dir,
+        intrinsic_path,
+        mode,
+        min_overlap_score=0.0,
+        pose_dir=None,
+    ):
+        augment_fn = self.augment_fn if mode == "train" else None
+        data_source = (
+            self.trainval_data_source
+            if mode in ["train", "val"]
+            else self.test_data_source
+        )
+        if str(data_source).lower() == "megadepth":
+            npz_names = [f"{n}.npz" for n in npz_names]
+        with tqdm_joblib(
+            tqdm(
+                desc=f"[rank:{self.rank}] loading {mode} datasets",
+                total=len(npz_names),
+                disable=int(self.rank) != 0,
+            )
+        ):
+            if data_source == "ScanNet":
+                size = self.scan_img_resize
+                fp = self.fp16
+                datasets = Parallel(
+                    n_jobs=math.floor(
+                        len(os.sched_getaffinity(0)) *
+                        0.9 / comm.get_local_size()
+                    )
+                )(
+                    delayed(
+                        lambda x: _build_dataset(
+                            ScanNetDataset,
+                            data_root,
+                            osp.join(npz_dir, x),
+                            intrinsic_path,
+                            mode=mode,
+                            min_overlap_score=min_overlap_score,
+                            augment_fn=augment_fn,
+                            pose_dir=pose_dir,
+                            img_resize=size,
+                            fp16=fp,
+                        )
+                    )(name)
+                    for name in npz_names
+                )
+            elif data_source == "MegaDepth":
+                # TODO: _pickle.PicklingError: Could not pickle the task to send it to the workers.
+                raise NotImplementedError()
+                datasets = Parallel(
+                    n_jobs=math.floor(
+                        len(os.sched_getaffinity(0)) *
+                        0.9 / comm.get_local_size()
+                    )
+                )(
+                    delayed(
+                        lambda x: _build_dataset(
+                            MegaDepthDataset,
+                            data_root,
+                            osp.join(npz_dir, x),
+                            mode=mode,
+                            min_overlap_score=min_overlap_score,
+                            img_resize=self.mgdpt_img_resize,
+                            df=self.mgdpt_df,
+                            img_padding=self.mgdpt_img_pad,
+                            depth_padding=self.mgdpt_depth_pad,
+                            augment_fn=augment_fn,
+                            coarse_scale=self.coarse_scale,
+                        )
+                    )(name)
+                    for name in npz_names
+                )
+            else:
+                raise ValueError(f"Unknown dataset: {data_source}")
+        return ConcatDataset(datasets)
+
+    def train_dataloader(self):
+        """Build training dataloader for ScanNet / MegaDepth."""
+        assert self.data_sampler in ["scene_balance"]
+        logger.info(
+            f"[rank:{self.rank}/{self.world_size}]: Train Sampler and DataLoader re-init (should not re-init between epochs!)."
+        )
+        if self.data_sampler == "scene_balance":
+            sampler = RandomConcatSampler(
+                self.train_dataset,
+                self.n_samples_per_subset,
+                self.subset_replacement,
+                self.shuffle,
+                self.repeat,
+                self.seed,
+            )
+        else:
+            sampler = None
+        dataloader = DataLoader(
+            self.train_dataset, sampler=sampler, **self.train_loader_params
+        )
+        return dataloader
+
+    def val_dataloader(self):
+        """Build validation dataloader for ScanNet / MegaDepth."""
+        logger.info(
+            f"[rank:{self.rank}/{self.world_size}]: Val Sampler and DataLoader re-init."
+        )
+        if not isinstance(self.val_dataset, abc.Sequence):
+            sampler = DistributedSampler(self.val_dataset, shuffle=False)
+            return DataLoader(
+                self.val_dataset, sampler=sampler, **self.val_loader_params
+            )
+        else:
+            dataloaders = []
+            for dataset in self.val_dataset:
+                sampler = DistributedSampler(dataset, shuffle=False)
+                dataloaders.append(
+                    DataLoader(dataset, sampler=sampler,
+                               **self.val_loader_params)
+                )
+            return dataloaders
+
+    def test_dataloader(self, *args, **kwargs):
+        logger.info(
+            f"[rank:{self.rank}/{self.world_size}]: Test Sampler and DataLoader re-init."
+        )
+        sampler = DistributedSampler(self.test_dataset, shuffle=False)
+        return DataLoader(self.test_dataset, sampler=sampler, **self.test_loader_params)
+
+
+def _build_dataset(dataset: Dataset, *args, **kwargs):
+    return dataset(*args, **kwargs)