Initial commit

ssd/modeling/box_head/__init__.py

+from ssd.modeling import registry
+from .box_head import SSDBoxHead
+
+__all__ = ['build_box_head', 'SSDBoxHead']
+
+
+def build_box_head(cfg):
+    return registry.BOX_HEADS[cfg.MODEL.BOX_HEAD.NAME](cfg)

ssd/modeling/box_head/box_head.py

+from torch import nn
+import torch.nn.functional as F
+
+from ssd.modeling import registry
+from ssd.modeling.anchors.prior_box import PriorBox
+from ssd.modeling.box_head.box_predictor import make_box_predictor
+from ssd.utils import box_utils
+from .inference import PostProcessor
+from .loss import MultiBoxLoss
+
+
+@registry.BOX_HEADS.register('SSDBoxHead')
+class SSDBoxHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.predictor = make_box_predictor(cfg)
+        self.loss_evaluator = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
+        self.post_processor = PostProcessor(cfg)
+        self.priors = None
+
+    def forward(self, features, targets=None):
+        cls_logits, bbox_pred = self.predictor(features)
+        if self.training:
+            return self._forward_train(cls_logits, bbox_pred, targets)
+        else:
+            return self._forward_test(cls_logits, bbox_pred)
+
+    def _forward_train(self, cls_logits, bbox_pred, targets):
+        gt_boxes, gt_labels = targets['boxes'], targets['labels']
+        reg_loss, cls_loss = self.loss_evaluator(cls_logits, bbox_pred, gt_labels, gt_boxes)
+        loss_dict = dict(
+            reg_loss=reg_loss,
+            cls_loss=cls_loss,
+        )
+        detections = (cls_logits, bbox_pred)
+        return detections, loss_dict
+
+    def _forward_test(self, cls_logits, bbox_pred):
+        if self.priors is None:
+            self.priors = PriorBox(self.cfg)().to(bbox_pred.device)
+        scores = F.softmax(cls_logits, dim=2)
+        boxes = box_utils.convert_locations_to_boxes(
+            bbox_pred, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE
+        )
+        boxes = box_utils.center_form_to_corner_form(boxes)
+        detections = (scores, boxes)
+        detections = self.post_processor(detections)
+        return detections, {}

ssd/modeling/box_head/box_predictor.py

+import torch
+from torch import nn
+
+from ssd.layers import SeparableConv2d
+from ssd.modeling import registry
+
+
+class BoxPredictor(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.cls_headers = nn.ModuleList()
+        self.reg_headers = nn.ModuleList()
+        for level, (boxes_per_location, out_channels) in enumerate(zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)):
+            self.cls_headers.append(self.cls_block(level, out_channels, boxes_per_location))
+            self.reg_headers.append(self.reg_block(level, out_channels, boxes_per_location))
+        self.reset_parameters()
+
+    def cls_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, features):
+        cls_logits = []
+        bbox_pred = []
+        for feature, cls_header, reg_header in zip(features, self.cls_headers, self.reg_headers):
+            cls_logits.append(cls_header(feature).permute(0, 2, 3, 1).contiguous())
+            bbox_pred.append(reg_header(feature).permute(0, 2, 3, 1).contiguous())
+
+        batch_size = features[0].shape[0]
+        cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1, self.cfg.MODEL.NUM_CLASSES)
+        bbox_pred = torch.cat([l.view(l.shape[0], -1) for l in bbox_pred], dim=1).view(batch_size, -1, 4)
+
+        return cls_logits, bbox_pred
+
+
+@registry.BOX_PREDICTORS.register('SSDBoxPredictor')
+class SSDBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+@registry.BOX_PREDICTORS.register('SSDLiteBoxPredictor')
+class SSDLiteBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+def make_box_predictor(cfg):
+    return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg)

ssd/modeling/box_head/inference.py

+import torch
+
+from ssd.structures.container import Container
+from ssd.utils.nms import batched_nms
+
+
+class PostProcessor:
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.width = cfg.INPUT.IMAGE_SIZE
+        self.height = cfg.INPUT.IMAGE_SIZE
+
+    def __call__(self, detections):
+        batches_scores, batches_boxes = detections
+        device = batches_scores.device
+        batch_size = batches_scores.size(0)
+        results = []
+        for batch_id in range(batch_size):
+            scores, boxes = batches_scores[batch_id], batches_boxes[batch_id]  # (N, #CLS) (N, 4)
+            num_boxes = scores.shape[0]
+            num_classes = scores.shape[1]
+
+            boxes = boxes.view(num_boxes, 1, 4).expand(num_boxes, num_classes, 4)
+            labels = torch.arange(num_classes, device=device)
+            labels = labels.view(1, num_classes).expand_as(scores)
+
+            # remove predictions with the background label
+            boxes = boxes[:, 1:]
+            scores = scores[:, 1:]
+            labels = labels[:, 1:]
+
+            # batch everything, by making every class prediction be a separate instance
+            boxes = boxes.reshape(-1, 4)
+            scores = scores.reshape(-1)
+            labels = labels.reshape(-1)
+
+            # remove low scoring boxes
+            indices = torch.nonzero(scores > self.cfg.TEST.CONFIDENCE_THRESHOLD).squeeze(1)
+            boxes, scores, labels = boxes[indices], scores[indices], labels[indices]
+
+            boxes[:, 0::2] *= self.width
+            boxes[:, 1::2] *= self.height
+
+            keep = batched_nms(boxes, scores, labels, self.cfg.TEST.NMS_THRESHOLD)
+            # keep only topk scoring predictions
+            keep = keep[:self.cfg.TEST.MAX_PER_IMAGE]
+            boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
+
+            container = Container(boxes=boxes, labels=labels, scores=scores)
+            container.img_width = self.width
+            container.img_height = self.height
+            results.append(container)
+        return results

ssd/modeling/box_head/loss.py

+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from ssd.utils import box_utils
+
+
+class MultiBoxLoss(nn.Module):
+    def __init__(self, neg_pos_ratio):
+        """Implement SSD MultiBox Loss.
+
+        Basically, MultiBox loss combines classification loss
+         and Smooth L1 regression loss.
+        """
+        super(MultiBoxLoss, self).__init__()
+        self.neg_pos_ratio = neg_pos_ratio
+
+    def forward(self, confidence, predicted_locations, labels, gt_locations):
+        """Compute classification loss and smooth l1 loss.
+
+        Args:
+            confidence (batch_size, num_priors, num_classes): class predictions.
+            predicted_locations (batch_size, num_priors, 4): predicted locations.
+            labels (batch_size, num_priors): real labels of all the priors.
+            gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors.
+        """
+        num_classes = confidence.size(2)
+        with torch.no_grad():
+            # derived from cross_entropy=sum(log(p))
+            loss = -F.log_softmax(confidence, dim=2)[:, :, 0]
+            mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio)
+
+        confidence = confidence[mask, :]
+        classification_loss = F.cross_entropy(confidence.view(-1, num_classes), labels[mask], reduction='sum')
+
+        pos_mask = labels > 0
+        predicted_locations = predicted_locations[pos_mask, :].view(-1, 4)
+        gt_locations = gt_locations[pos_mask, :].view(-1, 4)
+        smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum')
+        num_pos = gt_locations.size(0)
+        return smooth_l1_loss / num_pos, classification_loss / num_pos

ssd/modeling/detector/__init__.py

+from .ssd_detector import SSDDetector
+
+_DETECTION_META_ARCHITECTURES = {
+    "SSDDetector": SSDDetector
+}
+
+
+def build_detection_model(cfg):
+    meta_arch = _DETECTION_META_ARCHITECTURES[cfg.MODEL.META_ARCHITECTURE]
+    return meta_arch(cfg)

ssd/modeling/detector/ssd_detector.py

+from torch import nn
+
+from ssd.modeling.backbone import build_backbone
+from ssd.modeling.box_head import build_box_head
+
+
+class SSDDetector(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.backbone = build_backbone(cfg)
+        self.box_head = build_box_head(cfg)
+
+    def forward(self, images, targets=None):
+        features = self.backbone(images)
+        detections, detector_losses = self.box_head(features, targets)
+        if self.training:
+            return detector_losses
+        return detections

ssd/modeling/registry.py

+from ssd.utils.registry import Registry
+
+BACKBONES = Registry()
+BOX_HEADS = Registry()
+BOX_PREDICTORS = Registry()

ssd/solver/build.py

+import torch
+
+from .lr_scheduler import WarmupMultiStepLR
+
+
+def make_optimizer(cfg, model, lr=None):
+    lr = cfg.SOLVER.BASE_LR if lr is None else lr
+    return torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
+
+
+def make_lr_scheduler(cfg, optimizer, milestones=None):
+    return WarmupMultiStepLR(optimizer=optimizer,
+                             milestones=cfg.SOLVER.LR_STEPS if milestones is None else milestones,
+                             gamma=cfg.SOLVER.GAMMA,
+                             warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
+                             warmup_iters=cfg.SOLVER.WARMUP_ITERS)

ssd/solver/lr_scheduler.py

+from bisect import bisect_right
+
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class WarmupMultiStepLR(_LRScheduler):
+    def __init__(self, optimizer, milestones, gamma=0.1, warmup_factor=1.0 / 3,
+                 warmup_iters=500, last_epoch=-1):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                "Milestones should be a list of" " increasing integers. Got {}",
+                milestones,
+            )
+
+        self.milestones = milestones
+        self.gamma = gamma
+        self.warmup_factor = warmup_factor
+        self.warmup_iters = warmup_iters
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        warmup_factor = 1
+        if self.last_epoch < self.warmup_iters:
+            alpha = float(self.last_epoch) / self.warmup_iters
+            warmup_factor = self.warmup_factor * (1 - alpha) + alpha
+        return [
+            base_lr
+            * warmup_factor
+            * self.gamma ** bisect_right(self.milestones, self.last_epoch)
+            for base_lr in self.base_lrs
+        ]

ssd/structures/container.py

+class Container:
+    """
+    Help class for manage boxes, labels, etc...
+    Not inherit dict due to `default_collate` will change dict's subclass to dict.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._data_dict = dict(*args, **kwargs)
+
+    def __setattr__(self, key, value):
+        object.__setattr__(self, key, value)
+
+    def __getitem__(self, key):
+        return self._data_dict[key]
+
+    def __iter__(self):
+        return self._data_dict.__iter__()
+
+    def __setitem__(self, key, value):
+        self._data_dict[key] = value
+
+    def _call(self, name, *args, **kwargs):
+        keys = list(self._data_dict.keys())
+        for key in keys:
+            value = self._data_dict[key]
+            if hasattr(value, name):
+                self._data_dict[key] = getattr(value, name)(*args, **kwargs)
+        return self
+
+    def to(self, *args, **kwargs):
+        return self._call('to', *args, **kwargs)
+
+    def numpy(self):
+        return self._call('numpy')
+
+    def resize(self, size):
+        """resize boxes
+        Args:
+            size: (width, height)
+        Returns:
+            self
+        """
+        img_width = getattr(self, 'img_width', -1)
+        img_height = getattr(self, 'img_height', -1)
+        assert img_width > 0 and img_height > 0
+        assert 'boxes' in self._data_dict
+        boxes = self._data_dict['boxes']
+        new_width, new_height = size
+        boxes[:, 0::2] *= (new_width / img_width)
+        boxes[:, 1::2] *= (new_height / img_height)
+        return self
+
+    def __repr__(self):
+        return self._data_dict.__repr__()

ssd/utils/__init__.py

+from .misc import *

ssd/utils/box_utils.py

+import torch
+import math
+
+
+def convert_locations_to_boxes(locations, priors, center_variance,
+                               size_variance):
+    """Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w).
+
+    The conversion:
+        $$predicted\_center * center_variance = \frac {real\_center - prior\_center} {prior\_hw}$$
+        $$exp(predicted\_hw * size_variance) = \frac {real\_hw} {prior\_hw}$$
+    We do it in the inverse direction here.
+    Args:
+        locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well.
+        priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes.
+        center_variance: a float used to change the scale of center.
+        size_variance: a float used to change of scale of size.
+    Returns:
+        boxes:  priors: [[center_x, center_y, w, h]]. All the values
+            are relative to the image size.
+    """
+    # priors can have one dimension less.
+    if priors.dim() + 1 == locations.dim():
+        priors = priors.unsqueeze(0)
+    return torch.cat([
+        locations[..., :2] * center_variance * priors[..., 2:] + priors[..., :2],
+        torch.exp(locations[..., 2:] * size_variance) * priors[..., 2:]
+    ], dim=locations.dim() - 1)
+
+
+def convert_boxes_to_locations(center_form_boxes, center_form_priors, center_variance, size_variance):
+    # priors can have one dimension less
+    if center_form_priors.dim() + 1 == center_form_boxes.dim():
+        center_form_priors = center_form_priors.unsqueeze(0)
+    return torch.cat([
+        (center_form_boxes[..., :2] - center_form_priors[..., :2]) / center_form_priors[..., 2:] / center_variance,
+        torch.log(center_form_boxes[..., 2:] / center_form_priors[..., 2:]) / size_variance
+    ], dim=center_form_boxes.dim() - 1)
+
+
+def area_of(left_top, right_bottom) -> torch.Tensor:
+    """Compute the areas of rectangles given two corners.
+
+    Args:
+        left_top (N, 2): left top corner.
+        right_bottom (N, 2): right bottom corner.
+
+    Returns:
+        area (N): return the area.
+    """
+    hw = torch.clamp(right_bottom - left_top, min=0.0)
+    return hw[..., 0] * hw[..., 1]
+
+
+def iou_of(boxes0, boxes1, eps=1e-5):
+    """Return intersection-over-union (Jaccard index) of boxes.
+
+    Args:
+        boxes0 (N, 4): ground truth boxes.
+        boxes1 (N or 1, 4): predicted boxes.
+        eps: a small number to avoid 0 as denominator.
+    Returns:
+        iou (N): IoU values.
+    """
+    overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2])
+    overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:])
+
+    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
+    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
+    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
+    return overlap_area / (area0 + area1 - overlap_area + eps)
+
+
+def assign_priors(gt_boxes, gt_labels, corner_form_priors,
+                  iou_threshold):
+    """Assign ground truth boxes and targets to priors.
+
+    Args:
+        gt_boxes (num_targets, 4): ground truth boxes.
+        gt_labels (num_targets): labels of targets.
+        priors (num_priors, 4): corner form priors
+    Returns:
+        boxes (num_priors, 4): real values for priors.
+        labels (num_priros): labels for priors.
+    """
+    # size: num_priors x num_targets
+    ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1))
+    # size: num_priors
+    best_target_per_prior, best_target_per_prior_index = ious.max(1)
+    # size: num_targets
+    best_prior_per_target, best_prior_per_target_index = ious.max(0)
+
+    for target_index, prior_index in enumerate(best_prior_per_target_index):
+        best_target_per_prior_index[prior_index] = target_index
+    # 2.0 is used to make sure every target has a prior assigned
+    best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2)
+    # size: num_priors
+    labels = gt_labels[best_target_per_prior_index]
+    labels[best_target_per_prior < iou_threshold] = 0  # the backgournd id
+    boxes = gt_boxes[best_target_per_prior_index]
+    return boxes, labels
+
+
+def hard_negative_mining(loss, labels, neg_pos_ratio):
+    """
+    It used to suppress the presence of a large number of negative prediction.
+    It works on image level not batch level.
+    For any example/image, it keeps all the positive predictions and
+     cut the number of negative predictions to make sure the ratio
+     between the negative examples and positive examples is no more
+     the given ratio for an image.
+
+    Args:
+        loss (N, num_priors): the loss for each example.
+        labels (N, num_priors): the labels.
+        neg_pos_ratio:  the ratio between the negative examples and positive examples.
+    """
+    pos_mask = labels > 0
+    num_pos = pos_mask.long().sum(dim=1, keepdim=True)
+    num_neg = num_pos * neg_pos_ratio
+
+    loss[pos_mask] = -math.inf
+    _, indexes = loss.sort(dim=1, descending=True)
+    _, orders = indexes.sort(dim=1)
+    neg_mask = orders < num_neg
+    return pos_mask | neg_mask
+
+
+def center_form_to_corner_form(locations):
+    return torch.cat([locations[..., :2] - locations[..., 2:] / 2,
+                      locations[..., :2] + locations[..., 2:] / 2], locations.dim() - 1)
+
+
+def corner_form_to_center_form(boxes):
+    return torch.cat([
+        (boxes[..., :2] + boxes[..., 2:]) / 2,
+        boxes[..., 2:] - boxes[..., :2]
+    ], boxes.dim() - 1)

ssd/utils/checkpoint.py

+import logging
+import os
+
+import torch
+from torch.nn.parallel import DistributedDataParallel
+
+from ssd.utils.model_zoo import cache_url
+
+
+class CheckPointer:
+    _last_checkpoint_name = 'last_checkpoint.txt'
+
+    def __init__(self,
+                 model,
+                 optimizer=None,
+                 scheduler=None,
+                 save_dir="",
+                 save_to_disk=None,
+                 logger=None):
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.save_dir = save_dir
+        self.save_to_disk = save_to_disk
+        if logger is None:
+            logger = logging.getLogger(__name__)
+        self.logger = logger
+
+    def save(self, name, **kwargs):
+        if not self.save_dir:
+            return
+
+        if not self.save_to_disk:
+            return
+
+        data = {}
+        if isinstance(self.model, DistributedDataParallel):
+            data['model'] = self.model.module.state_dict()
+        else:
+            data['model'] = self.model.state_dict()
+        if self.optimizer is not None:
+            data["optimizer"] = self.optimizer.state_dict()
+        if self.scheduler is not None:
+            data["scheduler"] = self.scheduler.state_dict()
+        data.update(kwargs)
+
+        save_file = os.path.join(self.save_dir, "{}.pth".format(name))
+        self.logger.info("Saving checkpoint to {}".format(save_file))
+        torch.save(data, save_file)
+
+        self.tag_last_checkpoint(save_file)
+
+    def load(self, f=None, use_latest=True):
+        if self.has_checkpoint() and use_latest:
+            # override argument with existing checkpoint
+            f = self.get_checkpoint_file()
+        if not f:
+            # no checkpoint could be found
+            self.logger.info("No checkpoint found.")
+            return {}
+
+        self.logger.info("Loading checkpoint from {}".format(f))
+        checkpoint = self._load_file(f)
+        model = self.model
+        if isinstance(model, DistributedDataParallel):
+            model = self.model.module
+
+        model.load_state_dict(checkpoint.pop("model"))
+        if "optimizer" in checkpoint and self.optimizer:
+            self.logger.info("Loading optimizer from {}".format(f))
+            self.optimizer.load_state_dict(checkpoint.pop("optimizer"))
+        if "scheduler" in checkpoint and self.scheduler:
+            self.logger.info("Loading scheduler from {}".format(f))
+            self.scheduler.load_state_dict(checkpoint.pop("scheduler"))
+
+        # return any further checkpoint data
+        return checkpoint
+
+    def get_checkpoint_file(self):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        try:
+            with open(save_file, "r") as f:
+                last_saved = f.read()
+                last_saved = last_saved.strip()
+        except IOError:
+            # if file doesn't exist, maybe because it has just been
+            # deleted by a separate process
+            last_saved = ""
+        return last_saved
+
+    def has_checkpoint(self):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        return os.path.exists(save_file)
+
+    def tag_last_checkpoint(self, last_filename):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        with open(save_file, "w") as f:
+            f.write(last_filename)
+
+    def _load_file(self, f):
+        # download url files
+        if f.startswith("http"):
+            # if the file is a url path, download it and cache it
+            cached_f = cache_url(f)
+            self.logger.info("url {} cached in {}".format(f, cached_f))
+            f = cached_f
+        return torch.load(f, map_location=torch.device("cpu"))

ssd/utils/dist_util.py

+import pickle
+
+import torch
+import torch.distributed as dist
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    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()
+
+
+def _encode(encoded_data, data):
+    # gets a byte representation for the data
+    encoded_bytes = pickle.dumps(data)
+    # convert this byte string into a byte tensor
+    storage = torch.ByteStorage.from_buffer(encoded_bytes)
+    tensor = torch.ByteTensor(storage).to("cuda")
+    # encoding: first byte is the size and then rest is the data
+    s = tensor.numel()
+    assert s <= 255, "Can't encode data greater than 255 bytes"
+    # put the encoded data in encoded_data
+    encoded_data[0] = s
+    encoded_data[1: (s + 1)] = tensor
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.LongTensor([tensor.numel()]).to("cuda")
+    size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    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

ssd/utils/logger.py

+import logging
+import os
+import sys
+
+
+def setup_logger(name, distributed_rank, save_dir=None):
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    # don't log results for the non-master process
+    if distributed_rank > 0:
+        return logger
+    stream_handler = logging.StreamHandler(stream=sys.stdout)
+    stream_handler.setLevel(logging.DEBUG)
+    formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+    if save_dir:
+        fh = logging.FileHandler(os.path.join(save_dir, 'log.txt'))
+        fh.setLevel(logging.DEBUG)
+        fh.setFormatter(formatter)
+        logger.addHandler(fh)
+    return logger

ssd/utils/metric_logger.py

+from collections import deque, defaultdict
+import numpy as np
+import torch
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=10):
+        self.deque = deque(maxlen=window_size)
+        self.value = np.nan
+        self.series = []
+        self.total = 0.0
+        self.count = 0
+
+    def update(self, value):
+        self.deque.append(value)
+        self.series.append(value)
+        self.count += 1
+        self.total += value
+        self.value = value
+
+    @property
+    def median(self):
+        values = np.array(self.deque)
+        return np.median(values)
+
+    @property
+    def avg(self):
+        values = np.array(self.deque)
+        return np.mean(values)
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+
+class MetricLogger:
+    def __init__(self, delimiter=", "):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {:.3f} ({:.3f})".format(name, meter.avg, meter.global_avg)
+            )
+        return self.delimiter.join(loss_str)

ssd/utils/misc.py

+import errno
+import os
+
+
+def str2bool(s):
+    return s.lower() in ('true', '1')
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise