dnt 0.2.4__py3-none-any.whl → 0.3.1.7__py3-none-any.whl

dnt/third_party/fast-reid/fastreid/utils/comm.py

@@ -0,0 +1,255 @@
+"""
+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

dnt/third_party/fast-reid/fastreid/utils/compute_dist.py

@@ -0,0 +1,200 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+# Modified from: https://github.com/open-mmlab/OpenUnReID/blob/66bb2ae0b00575b80fbe8915f4d4f4739cc21206/openunreid/core/utils/compute_dist.py
+
+
+import faiss
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .faiss_utils import (
+    index_init_cpu,
+    index_init_gpu,
+    search_index_pytorch,
+    search_raw_array_pytorch,
+)
+
+__all__ = [
+    "build_dist",
+    "compute_jaccard_distance",
+    "compute_euclidean_distance",
+    "compute_cosine_distance",
+]
+
+
+@torch.no_grad()
+def build_dist(feat_1: torch.Tensor, feat_2: torch.Tensor, metric: str = "euclidean", **kwargs) -> np.ndarray:
+    r"""Compute distance between two feature embeddings.
+
+    Args:
+        feat_1 (torch.Tensor): 2-D feature with batch dimension.
+        feat_2 (torch.Tensor): 2-D feature with batch dimension.
+        metric:
+
+    Returns:
+        numpy.ndarray: distance matrix.
+    """
+    assert metric in ["cosine", "euclidean", "jaccard"], "Expected metrics are cosine, euclidean and jaccard, " \
+                                                         "but got {}".format(metric)
+
+    if metric == "euclidean":
+        return compute_euclidean_distance(feat_1, feat_2)
+
+    elif metric == "cosine":
+        return compute_cosine_distance(feat_1, feat_2)
+
+    elif metric == "jaccard":
+        feat = torch.cat((feat_1, feat_2), dim=0)
+        dist = compute_jaccard_distance(feat, k1=kwargs["k1"], k2=kwargs["k2"], search_option=0)
+        return dist[: feat_1.size(0), feat_1.size(0):]
+
+
+def k_reciprocal_neigh(initial_rank, i, k1):
+    forward_k_neigh_index = initial_rank[i, : k1 + 1]
+    backward_k_neigh_index = initial_rank[forward_k_neigh_index, : k1 + 1]
+    fi = np.where(backward_k_neigh_index == i)[0]
+    return forward_k_neigh_index[fi]
+
+
+@torch.no_grad()
+def compute_jaccard_distance(features, k1=20, k2=6, search_option=0, fp16=False):
+    if search_option < 3:
+        # torch.cuda.empty_cache()
+        features = features.cuda()
+
+    ngpus = faiss.get_num_gpus()
+    N = features.size(0)
+    mat_type = np.float16 if fp16 else np.float32
+
+    if search_option == 0:
+        # GPU + PyTorch CUDA Tensors (1)
+        res = faiss.StandardGpuResources()
+        res.setDefaultNullStreamAllDevices()
+        _, initial_rank = search_raw_array_pytorch(res, features, features, k1)
+        initial_rank = initial_rank.cpu().numpy()
+    elif search_option == 1:
+        # GPU + PyTorch CUDA Tensors (2)
+        res = faiss.StandardGpuResources()
+        index = faiss.GpuIndexFlatL2(res, features.size(-1))
+        index.add(features.cpu().numpy())
+        _, initial_rank = search_index_pytorch(index, features, k1)
+        res.syncDefaultStreamCurrentDevice()
+        initial_rank = initial_rank.cpu().numpy()
+    elif search_option == 2:
+        # GPU
+        index = index_init_gpu(ngpus, features.size(-1))
+        index.add(features.cpu().numpy())
+        _, initial_rank = index.search(features.cpu().numpy(), k1)
+    else:
+        # CPU
+        index = index_init_cpu(features.size(-1))
+        index.add(features.cpu().numpy())
+        _, initial_rank = index.search(features.cpu().numpy(), k1)
+
+    nn_k1 = []
+    nn_k1_half = []
+    for i in range(N):
+        nn_k1.append(k_reciprocal_neigh(initial_rank, i, k1))
+        nn_k1_half.append(k_reciprocal_neigh(initial_rank, i, int(np.around(k1 / 2))))
+
+    V = np.zeros((N, N), dtype=mat_type)
+    for i in range(N):
+        k_reciprocal_index = nn_k1[i]
+        k_reciprocal_expansion_index = k_reciprocal_index
+        for candidate in k_reciprocal_index:
+            candidate_k_reciprocal_index = nn_k1_half[candidate]
+            if len(
+                    np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)
+            ) > 2 / 3 * len(candidate_k_reciprocal_index):
+                k_reciprocal_expansion_index = np.append(
+                    k_reciprocal_expansion_index, candidate_k_reciprocal_index
+                )
+
+        k_reciprocal_expansion_index = np.unique(
+            k_reciprocal_expansion_index
+        )  # element-wise unique
+
+        x = features[i].unsqueeze(0).contiguous()
+        y = features[k_reciprocal_expansion_index]
+        m, n = x.size(0), y.size(0)
+        dist = (
+                torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n)
+                + torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+        )
+        dist.addmm_(x, y.t(), beta=1, alpha=-2)
+
+        if fp16:
+            V[i, k_reciprocal_expansion_index] = (
+                F.softmax(-dist, dim=1).view(-1).cpu().numpy().astype(mat_type)
+            )
+        else:
+            V[i, k_reciprocal_expansion_index] = (
+                F.softmax(-dist, dim=1).view(-1).cpu().numpy()
+            )
+
+    del nn_k1, nn_k1_half, x, y
+    features = features.cpu()
+
+    if k2 != 1:
+        V_qe = np.zeros_like(V, dtype=mat_type)
+        for i in range(N):
+            V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0)
+        V = V_qe
+        del V_qe
+
+    del initial_rank
+
+    invIndex = []
+    for i in range(N):
+        invIndex.append(np.where(V[:, i] != 0)[0])  # len(invIndex)=all_num
+
+    jaccard_dist = np.zeros((N, N), dtype=mat_type)
+    for i in range(N):
+        temp_min = np.zeros((1, N), dtype=mat_type)
+        indNonZero = np.where(V[i, :] != 0)[0]
+        indImages = [invIndex[ind] for ind in indNonZero]
+        for j in range(len(indNonZero)):
+            temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(
+                V[i, indNonZero[j]], V[indImages[j], indNonZero[j]]
+            )
+
+        jaccard_dist[i] = 1 - temp_min / (2 - temp_min)
+
+    del invIndex, V
+
+    pos_bool = jaccard_dist < 0
+    jaccard_dist[pos_bool] = 0.0
+
+    return jaccard_dist
+
+
+@torch.no_grad()
+def compute_euclidean_distance(features, others):
+    m, n = features.size(0), others.size(0)
+    dist_m = (
+            torch.pow(features, 2).sum(dim=1, keepdim=True).expand(m, n)
+            + torch.pow(others, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+    )
+    dist_m.addmm_(1, -2, features, others.t())
+
+    return dist_m.cpu().numpy()
+
+
+@torch.no_grad()
+def compute_cosine_distance(features, others):
+    """Computes cosine distance.
+    Args:
+        features (torch.Tensor): 2-D feature matrix.
+        others (torch.Tensor): 2-D feature matrix.
+    Returns:
+        torch.Tensor: distance matrix.
+    """
+    features = F.normalize(features, p=2, dim=1)
+    others = F.normalize(others, p=2, dim=1)
+    dist_m = 1 - torch.mm(features, others.t())
+    return dist_m.cpu().numpy()

dnt/third_party/fast-reid/fastreid/utils/env.py

@@ -0,0 +1,119 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import importlib
+import importlib.util
+import logging
+import numpy as np
+import os
+import random
+import sys
+from datetime import datetime
+import torch
+
+__all__ = ["seed_all_rng"]
+
+
+TORCH_VERSION = tuple(int(x) for x in torch.__version__.split(".")[:2])
+"""
+PyTorch version as a tuple of 2 ints. Useful for comparison.
+"""
+
+
+def seed_all_rng(seed=None):
+    """
+    Set the random seed for the RNG in torch, numpy and python.
+    Args:
+        seed (int): if None, will use a strong random seed.
+    """
+    if seed is None:
+        seed = (
+            os.getpid()
+            + int(datetime.now().strftime("%S%f"))
+            + int.from_bytes(os.urandom(2), "big")
+        )
+        logger = logging.getLogger(__name__)
+        logger.info("Using a generated random seed {}".format(seed))
+    np.random.seed(seed)
+    torch.set_rng_state(torch.manual_seed(seed).get_state())
+    random.seed(seed)
+
+
+# from https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path
+def _import_file(module_name, file_path, make_importable=False):
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    if make_importable:
+        sys.modules[module_name] = module
+    return module
+
+
+def _configure_libraries():
+    """
+    Configurations for some libraries.
+    """
+    # An environment option to disable `import cv2` globally,
+    # in case it leads to negative performance impact
+    disable_cv2 = int(os.environ.get("DETECTRON2_DISABLE_CV2", False))
+    if disable_cv2:
+        sys.modules["cv2"] = None
+    else:
+        # Disable opencl in opencv since its interaction with cuda often has negative effects
+        # This envvar is supported after OpenCV 3.4.0
+        os.environ["OPENCV_OPENCL_RUNTIME"] = "disabled"
+        try:
+            import cv2
+
+            if int(cv2.__version__.split(".")[0]) >= 3:
+                cv2.ocl.setUseOpenCL(False)
+        except ImportError:
+            pass
+
+    def get_version(module, digit=2):
+        return tuple(map(int, module.__version__.split(".")[:digit]))
+
+    # fmt: off
+    assert get_version(torch) >= (1, 4), "Requires torch>=1.4"
+    import yaml
+    assert get_version(yaml) >= (5, 1), "Requires pyyaml>=5.1"
+    # fmt: on
+
+
+_ENV_SETUP_DONE = False
+
+
+def setup_environment():
+    """Perform environment setup work. The default setup is a no-op, but this
+    function allows the user to specify a Python source file or a module in
+    the $FASTREID_ENV_MODULE environment variable, that performs
+    custom setup work that may be necessary to their computing environment.
+    """
+    global _ENV_SETUP_DONE
+    if _ENV_SETUP_DONE:
+        return
+    _ENV_SETUP_DONE = True
+
+    _configure_libraries()
+
+    custom_module_path = os.environ.get("FASTREID_ENV_MODULE")
+
+    if custom_module_path:
+        setup_custom_environment(custom_module_path)
+    else:
+        # The default setup is a no-op
+        pass
+
+
+def setup_custom_environment(custom_module):
+    """
+    Load custom environment setup by importing a Python source file or a
+    module, and run the setup function.
+    """
+    if custom_module.endswith(".py"):
+        module = _import_file("fastreid.utils.env.custom_module", custom_module)
+    else:
+        module = importlib.import_module(custom_module)
+    assert hasattr(module, "setup_environment") and callable(module.setup_environment), (
+        "Custom environment module defined in {} does not have the "
+        "required callable attribute 'setup_environment'."
+    ).format(custom_module)
+    module.setup_environment()