vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/TopicFM/third_party/matchformer/model/datasets/sampler.py

@@ -0,0 +1,77 @@
+import torch
+from torch.utils.data import Sampler, ConcatDataset
+
+
+class RandomConcatSampler(Sampler):
+    """ Random sampler for ConcatDataset. At each epoch, `n_samples_per_subset` samples will be draw from each subset
+    in the ConcatDataset. If `subset_replacement` is ``True``, sampling within each subset will be done with replacement.
+    However, it is impossible to sample data without replacement between epochs, unless bulding a stateful sampler lived along the entire training phase.
+    
+    For current implementation, the randomness of sampling is ensured no matter the sampler is recreated across epochs or not and call `torch.manual_seed()` or not.
+    Args:
+        shuffle (bool): shuffle the random sampled indices across all sub-datsets.
+        repeat (int): repeatedly use the sampled indices multiple times for training.
+            [arXiv:1902.05509, arXiv:1901.09335]
+    NOTE: Don't re-initialize the sampler between epochs (will lead to repeated samples)
+    NOTE: This sampler behaves differently with DistributedSampler.
+          It assume the dataset is splitted across ranks instead of replicated.
+    TODO: Add a `set_epoch()` method to fullfill sampling without replacement across epochs.
+          ref: https://github.com/PyTorchLightning/pytorch-lightning/blob/e9846dd758cfb1500eb9dba2d86f6912eb487587/pytorch_lightning/trainer/training_loop.py#L373
+    """
+    def __init__(self,
+                 data_source: ConcatDataset,
+                 n_samples_per_subset: int,
+                 subset_replacement: bool=True,
+                 shuffle: bool=True,
+                 repeat: int=1,
+                 seed: int=None):
+        if not isinstance(data_source, ConcatDataset):
+            raise TypeError("data_source should be torch.utils.data.ConcatDataset")
+        
+        self.data_source = data_source
+        self.n_subset = len(self.data_source.datasets)
+        self.n_samples_per_subset = n_samples_per_subset
+        self.n_samples = self.n_subset * self.n_samples_per_subset * repeat
+        self.subset_replacement = subset_replacement
+        self.repeat = repeat
+        self.shuffle = shuffle
+        self.generator = torch.manual_seed(seed)
+        assert self.repeat >= 1
+        
+    def __len__(self):
+        return self.n_samples
+    
+    def __iter__(self):
+        indices = []
+        # sample from each sub-dataset
+        for d_idx in range(self.n_subset):
+            low = 0 if d_idx==0 else self.data_source.cumulative_sizes[d_idx-1]
+            high = self.data_source.cumulative_sizes[d_idx]
+            if self.subset_replacement:
+                rand_tensor = torch.randint(low, high, (self.n_samples_per_subset, ),
+                                            generator=self.generator, dtype=torch.int64)
+            else:  # sample without replacement
+                len_subset = len(self.data_source.datasets[d_idx])
+                rand_tensor = torch.randperm(len_subset, generator=self.generator) + low
+                if len_subset >= self.n_samples_per_subset:
+                    rand_tensor = rand_tensor[:self.n_samples_per_subset]
+                else: # padding with replacement
+                    rand_tensor_replacement = torch.randint(low, high, (self.n_samples_per_subset - len_subset, ),
+                                                            generator=self.generator, dtype=torch.int64)
+                    rand_tensor = torch.cat([rand_tensor, rand_tensor_replacement])
+            indices.append(rand_tensor)
+        indices = torch.cat(indices)
+        if self.shuffle:  # shuffle the sampled dataset (from multiple subsets)
+            rand_tensor = torch.randperm(len(indices), generator=self.generator)
+            indices = indices[rand_tensor]
+
+        # repeat the sampled indices (can be used for RepeatAugmentation or pure RepeatSampling)
+        if self.repeat > 1:
+            repeat_indices = [indices.clone() for _ in range(self.repeat - 1)]
+            if self.shuffle:
+                _choice = lambda x: x[torch.randperm(len(x), generator=self.generator)]
+                repeat_indices = map(_choice, repeat_indices)
+            indices = torch.cat([indices, *repeat_indices], 0)
+        
+        assert indices.shape[0] == self.n_samples
+        return iter(indices.tolist())

vismatch/third_party/TopicFM/third_party/matchformer/model/datasets/scannet.py

@@ -0,0 +1,113 @@
+from os import path as osp
+from typing import Dict
+from unicodedata import name
+
+import numpy as np
+import torch
+import torch.utils as utils
+from numpy.linalg import inv
+from .dataset import (
+    read_scannet_gray,
+    read_scannet_depth,
+    read_scannet_pose
+)
+
+
+class ScanNetDataset(utils.data.Dataset):
+    def __init__(self,
+                 root_dir,
+                 npz_path,
+                 intrinsic_path,
+                 mode='train',
+                 min_overlap_score=0.4,
+                 augment_fn=None,
+                 pose_dir=None,
+                 **kwargs):
+        """Manage one scene of ScanNet Dataset.
+        Args:
+            root_dir (str): ScanNet root directory that contains scene folders.
+            npz_path (str): {scene_id}.npz path. This contains image pair information of a scene.
+            intrinsic_path (str): path to depth-camera intrinsic file.
+            mode (str): options are ['train', 'val', 'test'].
+            augment_fn (callable, optional): augments images with pre-defined visual effects.
+            pose_dir (str): ScanNet root directory that contains all poses.
+                (we use a separate (optional) pose_dir since we store images and poses separately.)
+        """
+        super().__init__()
+        self.root_dir = root_dir
+        self.pose_dir = pose_dir if pose_dir is not None else root_dir
+        self.mode = mode
+
+        # prepare data_names, intrinsics and extrinsics(T)
+        with np.load(npz_path) as data:
+            self.data_names = data['name']
+            if 'score' in data.keys() and mode not in ['val' or 'test']:
+                kept_mask = data['score'] > min_overlap_score
+                self.data_names = self.data_names[kept_mask]
+        self.intrinsics = dict(np.load(intrinsic_path))
+
+        # for training LoFTR
+        self.augment_fn = augment_fn if mode == 'train' else None
+
+    def __len__(self):
+        return len(self.data_names)
+
+    def _read_abs_pose(self, scene_name, name):
+        pth = osp.join(self.pose_dir,
+                       scene_name,
+                       'pose', f'{name}.txt')
+        return read_scannet_pose(pth)
+
+    def _compute_rel_pose(self, scene_name, name0, name1):
+        pose0 = self._read_abs_pose(scene_name, name0)
+        pose1 = self._read_abs_pose(scene_name, name1)
+        
+        return np.matmul(pose1, inv(pose0))  # (4, 4)
+
+    def __getitem__(self, idx):
+        data_name = self.data_names[idx]
+        scene_name, scene_sub_name, stem_name_0, stem_name_1 = data_name
+        scene_name = f'scene{scene_name:04d}_{scene_sub_name:02d}'
+
+        # read the grayscale image which will be resized to (1, 480, 640)
+        img_name0 = osp.join(self.root_dir, scene_name, 'color', f'{stem_name_0}.jpg')
+        img_name1 = osp.join(self.root_dir, scene_name, 'color', f'{stem_name_1}.jpg')
+        
+        # TODO: Support augmentation & handle seeds for each worker correctly.
+        image0 = read_scannet_gray(img_name0, resize=(640, 480), augment_fn=None)
+                                #    augment_fn=np.random.choice([self.augment_fn, None], p=[0.5, 0.5]))
+        image1 = read_scannet_gray(img_name1, resize=(640, 480), augment_fn=None)
+                                #    augment_fn=np.random.choice([self.augment_fn, None], p=[0.5, 0.5]))
+
+        # read the depthmap which is stored as (480, 640)
+        if self.mode in ['train', 'val']:
+            depth0 = read_scannet_depth(osp.join(self.root_dir, scene_name, 'depth', f'{stem_name_0}.png'))
+            depth1 = read_scannet_depth(osp.join(self.root_dir, scene_name, 'depth', f'{stem_name_1}.png'))
+        else:
+            depth0 = depth1 = torch.tensor([])
+
+        # read the intrinsic of depthmap
+        K_0 = K_1 = torch.tensor(self.intrinsics[scene_name].copy(), dtype=torch.float).reshape(3, 3)
+
+        # read and compute relative poses
+        T_0to1 = torch.tensor(self._compute_rel_pose(scene_name, stem_name_0, stem_name_1),
+                              dtype=torch.float32)
+        T_1to0 = T_0to1.inverse()
+
+        data = {
+            'image0': image0,   # (1, h, w)
+            'depth0': depth0,   # (h, w)
+            'image1': image1,
+            'depth1': depth1,
+            'T_0to1': T_0to1,   # (4, 4)
+            'T_1to0': T_1to0,
+            'K0': K_0,  # (3, 3)
+            'K1': K_1,
+            'dataset_name': 'ScanNet',
+            'scene_id': scene_name,
+            'pair_id': idx,
+            'pair_names': (osp.join(scene_name, 'color', f'{stem_name_0}.jpg'),
+                           osp.join(scene_name, 'color', f'{stem_name_1}.jpg'))
+        }
+
+        return data

vismatch/third_party/TopicFM/third_party/matchformer/model/lightning_loftr.py

@@ -0,0 +1,102 @@
+import pprint
+from loguru import logger
+from pathlib import Path
+
+import torch
+import numpy as np
+import pytorch_lightning as pl
+
+from .matchformer import Matchformer
+from .utils.metrics import (
+    compute_symmetrical_epipolar_errors,
+    compute_pose_errors,
+    aggregate_metrics
+)
+from .utils.comm import gather
+from .utils.misc import lower_config, flattenList
+from .utils.profiler import PassThroughProfiler
+
+
+class PL_LoFTR(pl.LightningModule):
+    def __init__(self, config, pretrained_ckpt=None, profiler=None, dump_dir=None):
+        super().__init__()
+        # Misc
+        self.config = config  # full config
+        _config = lower_config(self.config)
+        self.profiler = profiler or PassThroughProfiler()
+
+        # Matcher: LoFTR
+        self.matcher = Matchformer(config=_config['matchformer'])
+
+        # Pretrained weights
+        if pretrained_ckpt:
+            self.matcher.load_state_dict({k.replace('matcher.',''):v  for k,v in torch.load(pretrained_ckpt, map_location='cpu').items()})
+            logger.info(f"Load \'{pretrained_ckpt}\' as pretrained checkpoint")
+        
+        # Testing
+        self.dump_dir = dump_dir
+        
+    
+    def _compute_metrics(self, batch):
+        with self.profiler.profile("Copmute metrics"):
+            compute_symmetrical_epipolar_errors(batch)  # compute epi_errs for each match
+            compute_pose_errors(batch, self.config)  # compute R_errs, t_errs, pose_errs for each pair
+
+            rel_pair_names = list(zip(*batch['pair_names']))
+            bs = batch['image0'].size(0)
+            metrics = {
+                # to filter duplicate pairs caused by DistributedSampler
+                'identifiers': ['#'.join(rel_pair_names[b]) for b in range(bs)],
+                'epi_errs': [batch['epi_errs'][batch['m_bids'] == b].cpu().numpy() for b in range(bs)],
+                'R_errs': batch['R_errs'],
+                't_errs': batch['t_errs'],
+                'inliers': batch['inliers']}
+            ret_dict = {'metrics': metrics}
+        return ret_dict, rel_pair_names
+    
+
+    def test_step(self, batch, batch_idx):
+        with self.profiler.profile("LoFTR"):
+            self.matcher(batch)
+
+        ret_dict, rel_pair_names = self._compute_metrics(batch)
+
+        with self.profiler.profile("dump_results"):
+            if self.dump_dir is not None:
+                # dump results for further analysis
+                keys_to_save = {'mkpts0_f', 'mkpts1_f', 'mconf', 'epi_errs'}
+                pair_names = list(zip(*batch['pair_names']))
+                bs = batch['image0'].shape[0]
+                dumps = []
+                for b_id in range(bs):
+                    item = {}
+                    mask = batch['m_bids'] == b_id
+                    item['pair_names'] = pair_names[b_id]
+                    item['identifier'] = '#'.join(rel_pair_names[b_id])
+                    for key in keys_to_save:
+                        item[key] = batch[key][mask].cpu().numpy()
+                    for key in ['R_errs', 't_errs', 'inliers']:
+                        item[key] = batch[key][b_id]
+                    dumps.append(item)
+                ret_dict['dumps'] = dumps
+
+        return ret_dict
+
+    def test_epoch_end(self, outputs):
+        # metrics: dict of list, numpy
+        _metrics = [o['metrics'] for o in outputs]
+        metrics = {k: flattenList(gather(flattenList([_me[k] for _me in _metrics]))) for k in _metrics[0]}
+
+        # [{key: [{...}, *#bs]}, *#batch]
+        if self.dump_dir is not None:
+            Path(self.dump_dir).mkdir(parents=True, exist_ok=True)
+            _dumps = flattenList([o['dumps'] for o in outputs])  # [{...}, #bs*#batch]
+            dumps = flattenList(gather(_dumps))  # [{...}, #proc*#bs*#batch]
+            logger.info(f'Prediction and evaluation results will be saved to: {self.dump_dir}')
+
+        if self.trainer.global_rank == 0:
+            print(self.profiler.summary())
+            val_metrics_4tb = aggregate_metrics(metrics, self.config.TRAINER.EPI_ERR_THR)
+            logger.info('\n' + pprint.pformat(val_metrics_4tb))
+            if self.dump_dir is not None:
+                np.save(Path(self.dump_dir) / 'LoFTR_pred_eval', dumps)

vismatch/third_party/TopicFM/third_party/matchformer/model/matchformer.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+from .backbone import build_backbone
+from .backbone.fine_preprocess import FinePreprocess
+from .backbone.coarse_matching import CoarseMatching
+from .backbone.fine_matching import FineMatching
+from einops.einops import rearrange
+
+
+class Matchformer(nn.Module):
+    
+    def __init__(self, config):
+        super().__init__()
+        # Misc
+        self.config = config
+        self.backbone = build_backbone(config)
+        self.coarse_matching = CoarseMatching(config['match_coarse'])
+        self.fine_preprocess = FinePreprocess(config)
+        self.fine_matching = FineMatching()
+    
+    
+    def forward(self, data):
+        data.update({
+            'bs': data['image0'].size(0),
+            'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
+        })
+
+        mask_c0 = mask_c1 = None  # mask is useful in training
+        if 'mask0' in data:
+            mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)
+            
+        if data['hw0_i'] == data['hw1_i']:
+            feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
+            (feat_c0, feat_c1),(feat_f0, feat_f1) = feats_c.split(data['bs']),feats_f.split(data['bs'])
+        else:
+            (feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])
+
+        data.update({
+            'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
+            'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
+        })
+        
+        feat_c0 = rearrange(feat_c0, 'n c h w -> n (h w) c')
+        feat_c1 = rearrange(feat_c1, 'n c h w -> n (h w) c')
+
+        # match coarse-level
+        self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1)
+
+        # fine-level refinement
+        feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
+
+        # match fine-level
+        self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)
+        

vismatch/third_party/TopicFM/third_party/matchformer/model/utils/augment.py

@@ -0,0 +1,55 @@
+import albumentations as A
+
+
+class DarkAug(object):
+    """
+    Extreme dark augmentation aiming at Aachen Day-Night
+    """
+
+    def __init__(self) -> None:
+        self.augmentor = A.Compose([
+            A.RandomBrightnessContrast(p=0.75, brightness_limit=(-0.6, 0.0), contrast_limit=(-0.5, 0.3)),
+            A.Blur(p=0.1, blur_limit=(3, 9)),
+            A.MotionBlur(p=0.2, blur_limit=(3, 25)),
+            A.RandomGamma(p=0.1, gamma_limit=(15, 65)),
+            A.HueSaturationValue(p=0.1, val_shift_limit=(-100, -40))
+        ], p=0.75)
+
+    def __call__(self, x):
+        return self.augmentor(image=x)['image']
+
+
+class MobileAug(object):
+    """
+    Random augmentations aiming at images of mobile/handhold devices.
+    """
+
+    def __init__(self):
+        self.augmentor = A.Compose([
+            A.MotionBlur(p=0.25),
+            A.ColorJitter(p=0.5),
+            A.RandomRain(p=0.1),  # random occlusion
+            A.RandomSunFlare(p=0.1),
+            A.JpegCompression(p=0.25),
+            A.ISONoise(p=0.25)
+        ], p=1.0)
+
+    def __call__(self, x):
+        return self.augmentor(image=x)['image']
+
+
+def build_augmentor(method=None, **kwargs):
+    if method is not None:
+        raise NotImplementedError('Using of augmentation functions are not supported yet!')
+    if method == 'dark':
+        return DarkAug()
+    elif method == 'mobile':
+        return MobileAug()
+    elif method is None:
+        return None
+    else:
+        raise ValueError(f'Invalid augmentation method: {method}')
+
+
+if __name__ == '__main__':
+    augmentor = build_augmentor('FDA')

vismatch/third_party/TopicFM/third_party/matchformer/model/utils/comm.py

@@ -0,0 +1,265 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+[Copied from detectron2]
+This file contains primitives for multi-gpu communication.
+This is useful when doing distributed training.
+"""
+
+import functools
+import logging
+import numpy as np
+import pickle
+import torch
+import torch.distributed as dist
+
+_LOCAL_PROCESS_GROUP = None
+"""
+A torch process group which only includes processes that on the same machine as the current process.
+This variable is set when processes are spawned by `launch()` in "engine/launch.py".
+"""
+
+
+def get_world_size() -> int:
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank() -> int:
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def get_local_rank() -> int:
+    """
+    Returns:
+        The rank of the current process within the local (per-machine) process group.
+    """
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    assert _LOCAL_PROCESS_GROUP is not None
+    return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
+
+
+def get_local_size() -> int:
+    """
+    Returns:
+        The size of the per-machine process group,
+        i.e. the number of processes per machine.
+    """
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)
+
+
+def is_main_process() -> bool:
+    return get_rank() == 0
+
+
+def synchronize():
+    """
+    Helper function to synchronize (barrier) among all processes when
+    using distributed training
+    """
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        return
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return
+    dist.barrier()
+
+
+@functools.lru_cache()
+def _get_global_gloo_group():
+    """
+    Return a process group based on gloo backend, containing all the ranks
+    The result is cached.
+    """
+    if dist.get_backend() == "nccl":
+        return dist.new_group(backend="gloo")
+    else:
+        return dist.group.WORLD
+
+
+def _serialize_to_tensor(data, group):
+    backend = dist.get_backend(group)
+    assert backend in ["gloo", "nccl"]
+    device = torch.device("cpu" if backend == "gloo" else "cuda")
+
+    buffer = pickle.dumps(data)
+    if len(buffer) > 1024 ** 3:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
+                get_rank(), len(buffer) / (1024 ** 3), device
+            )
+        )
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to(device=device)
+    return tensor
+
+
+def _pad_to_largest_tensor(tensor, group):
+    """
+    Returns:
+        list[int]: size of the tensor, on each rank
+        Tensor: padded tensor that has the max size
+    """
+    world_size = dist.get_world_size(group=group)
+    assert (
+        world_size >= 1
+    ), "comm.gather/all_gather must be called from ranks within the given group!"
+    local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)
+    size_list = [
+        torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size)
+    ]
+    dist.all_gather(size_list, local_size, group=group)
+
+    size_list = [int(size.item()) for size in size_list]
+
+    max_size = max(size_list)
+
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    if local_size != max_size:
+        padding = torch.zeros((max_size - local_size,), dtype=torch.uint8, device=tensor.device)
+        tensor = torch.cat((tensor, padding), dim=0)
+    return size_list, tensor
+
+
+def all_gather(data, group=None):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors).
+
+    Args:
+        data: any picklable object
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    if get_world_size() == 1:
+        return [data]
+    if group is None:
+        group = _get_global_gloo_group()
+    if dist.get_world_size(group) == 1:
+        return [data]
+
+    tensor = _serialize_to_tensor(data, group)
+
+    size_list, tensor = _pad_to_largest_tensor(tensor, group)
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    tensor_list = [
+        torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
+    ]
+    dist.all_gather(tensor_list, tensor, group=group)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def gather(data, dst=0, group=None):
+    """
+    Run gather on arbitrary picklable data (not necessarily tensors).
+
+    Args:
+        data: any picklable object
+        dst (int): destination rank
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+
+    Returns:
+        list[data]: on dst, a list of data gathered from each rank. Otherwise,
+            an empty list.
+    """
+    if get_world_size() == 1:
+        return [data]
+    if group is None:
+        group = _get_global_gloo_group()
+    if dist.get_world_size(group=group) == 1:
+        return [data]
+    rank = dist.get_rank(group=group)
+
+    tensor = _serialize_to_tensor(data, group)
+    size_list, tensor = _pad_to_largest_tensor(tensor, group)
+
+    # receiving Tensor from all ranks
+    if rank == dst:
+        max_size = max(size_list)
+        tensor_list = [
+            torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
+        ]
+        dist.gather(tensor, tensor_list, dst=dst, group=group)
+
+        data_list = []
+        for size, tensor in zip(size_list, tensor_list):
+            buffer = tensor.cpu().numpy().tobytes()[:size]
+            data_list.append(pickle.loads(buffer))
+        return data_list
+    else:
+        dist.gather(tensor, [], dst=dst, group=group)
+        return []
+
+
+def shared_random_seed():
+    """
+    Returns:
+        int: a random number that is the same across all workers.
+            If workers need a shared RNG, they can use this shared seed to
+            create one.
+
+    All workers must call this function, otherwise it will deadlock.
+    """
+    ints = np.random.randint(2 ** 31)
+    all_ints = all_gather(ints)
+    return all_ints[0]
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Reduce the values in the dictionary from all processes so that process with rank
+    0 has the reduced results.
+
+    Args:
+        input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
+        average (bool): whether to do average or sum
+
+    Returns:
+        a dict with the same keys as input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.reduce(values, dst=0)
+        if dist.get_rank() == 0 and average:
+            # only main process gets accumulated, so only divide by
+            # world_size in this case
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict

vismatch/third_party/TopicFM/third_party/matchformer/model/utils/dataloader.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+
+# --- PL-DATAMODULE ---
+
+def get_local_split(items: list, world_size: int, rank: int, seed: int):
+    """ The local rank only loads a split of the dataset. """
+    n_items = len(items)
+    items_permute = np.random.RandomState(seed).permutation(items)
+    if n_items % world_size == 0:
+        padded_items = items_permute
+    else:
+        padding = np.random.RandomState(seed).choice(
+            items,
+            world_size - (n_items % world_size),
+            replace=True)
+        padded_items = np.concatenate([items_permute, padding])
+        assert len(padded_items) % world_size == 0, \
+            f'len(padded_items): {len(padded_items)}; world_size: {world_size}; len(padding): {len(padding)}'
+    n_per_rank = len(padded_items) // world_size
+    local_items = padded_items[n_per_rank * rank: n_per_rank * (rank+1)]
+
+    return local_items