From d426692b95c81a610d4eacab4a5ebbeee1932384 Mon Sep 17 00:00:00 2001
From: shaoniangu <shaoniangu@163.com>
Date: Fri, 5 Jul 2019 01:09:30 +0800
Subject: [PATCH] remove code that is irrelevant to the paper

---
 layers/__init__.py     |  77 ++----------
 layers/cluster_loss.py | 269 -----------------------------------------
 layers/range_loss.py   | 232 -----------------------------------
 3 files changed, 10 insertions(+), 568 deletions(-)
 delete mode 100644 layers/cluster_loss.py
 delete mode 100644 layers/range_loss.py

diff --git a/layers/__init__.py b/layers/__init__.py
index aa3ebcb..378af1e 100644
--- a/layers/__init__.py
+++ b/layers/__init__.py
@@ -7,22 +7,15 @@
 import torch.nn.functional as F
 
 from .triplet_loss import TripletLoss, CrossEntropyLabelSmooth
-from .cluster_loss import ClusterLoss
 from .center_loss import CenterLoss
-from .range_loss import RangeLoss
 
 
 def make_loss(cfg, num_classes):    # modified by gu
     sampler = cfg.DATALOADER.SAMPLER
     if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
         triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
-    elif cfg.MODEL.METRIC_LOSS_TYPE == 'cluster':
-        cluster = ClusterLoss(cfg.SOLVER.CLUSTER_MARGIN, True, True, cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE, cfg.DATALOADER.NUM_INSTANCE)
-    elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_cluster':
-        triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
-        cluster = ClusterLoss(cfg.SOLVER.CLUSTER_MARGIN, True, True, cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE, cfg.DATALOADER.NUM_INSTANCE)
     else:
-        print('expected METRIC_LOSS_TYPE should be triplet, cluster, triplet_cluster'
+        print('expected METRIC_LOSS_TYPE should be triplet'
               'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
 
     if cfg.MODEL.IF_LABELSMOOTH == 'on':
@@ -39,23 +32,11 @@ def make_loss(cfg, num_classes):    # modified by gu
         def loss_func(score, feat, target):
             if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
                 if cfg.MODEL.IF_LABELSMOOTH == 'on':
-                    return xent(score, target) + triplet(feat, target)[0]  # new add by luo, open label smooth
+                    return xent(score, target) + triplet(feat, target)[0]
                 else:
-                    return F.cross_entropy(score, target) + triplet(feat, target)[0]    # new add by luo, no label smooth
-
-            elif cfg.MODEL.METRIC_LOSS_TYPE == 'cluster':
-                if cfg.MODEL.IF_LABELSMOOTH == 'on':
-                    return xent(score, target) + cluster(feat, target)[0]  # new add by luo, open label smooth
-                else:
-                    return F.cross_entropy(score, target) + cluster(feat, target)[0]    # new add by luo, no label smooth
-
-            elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_cluster':
-                if cfg.MODEL.IF_LABELSMOOTH == 'on':
-                    return xent(score, target) + triplet(feat, target)[0] + cluster(feat, target)[0]  # new add by luo, open label smooth
-                else:
-                    return F.cross_entropy(score, target) + triplet(feat, target)[0] + cluster(feat, target)[0]    # new add by luo, no label smooth
+                    return F.cross_entropy(score, target) + triplet(feat, target)[0]
             else:
-                print('expected METRIC_LOSS_TYPE should be triplet, cluster, triplet_cluster,'
+                print('expected METRIC_LOSS_TYPE should be triplet'
                       'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
     else:
         print('expected sampler should be softmax, triplet or softmax_triplet, '
@@ -72,27 +53,12 @@ def make_loss_with_center(cfg, num_classes):    # modified by gu
     if cfg.MODEL.METRIC_LOSS_TYPE == 'center':
         center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center loss
 
-    elif cfg.MODEL.METRIC_LOSS_TYPE == 'range_center':
-        center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center_range loss
-        range_criterion = RangeLoss(k=cfg.SOLVER.RANGE_K, margin=cfg.SOLVER.RANGE_MARGIN, alpha=cfg.SOLVER.RANGE_ALPHA,
-                                    beta=cfg.SOLVER.RANGE_BETA, ordered=True, use_gpu=True,
-                                    ids_per_batch=cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE,
-                                    imgs_per_id=cfg.DATALOADER.NUM_INSTANCE)
-
     elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_center':
         triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
         center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center loss
 
-    elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_range_center':
-        triplet = TripletLoss(cfg.SOLVER.MARGIN)  # triplet loss
-        center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True)  # center_range loss
-        range_criterion = RangeLoss(k=cfg.SOLVER.RANGE_K, margin=cfg.SOLVER.RANGE_MARGIN, alpha=cfg.SOLVER.RANGE_ALPHA,
-                                    beta=cfg.SOLVER.RANGE_BETA, ordered=True, use_gpu=True,
-                                    ids_per_batch=cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE,
-                                    imgs_per_id=cfg.DATALOADER.NUM_INSTANCE)
     else:
-        print('expected METRIC_LOSS_TYPE with center should be center, '
-              'range_center,triplet_center, triplet_range_center '
+        print('expected METRIC_LOSS_TYPE with center should be center, triplet_center'
               'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
 
     if cfg.MODEL.IF_LABELSMOOTH == 'on':
@@ -103,45 +69,22 @@ def make_loss_with_center(cfg, num_classes):    # modified by gu
         if cfg.MODEL.METRIC_LOSS_TYPE == 'center':
             if cfg.MODEL.IF_LABELSMOOTH == 'on':
                 return xent(score, target) + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)  # new add by luo, open label smooth
+                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
             else:
                 return F.cross_entropy(score, target) + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)    # new add by luo, no label smooth
-
-        elif cfg.MODEL.METRIC_LOSS_TYPE == 'range_center':
-            if cfg.MODEL.IF_LABELSMOOTH == 'on':
-                return xent(score, target) + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
-                        cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0] # new add by luo, open label smooth
-            else:
-                return F.cross_entropy(score, target) + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
-                        cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0]     # new add by luo, no label smooth
+                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
 
         elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_center':
             if cfg.MODEL.IF_LABELSMOOTH == 'on':
                 return xent(score, target) + \
                         triplet(feat, target)[0] + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)  # new add by luo, open label smooth
+                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
             else:
                 return F.cross_entropy(score, target) + \
                         triplet(feat, target)[0] + \
-                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)    # new add by luo, no label smooth
-
-        elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_range_center':
-            if cfg.MODEL.IF_LABELSMOOTH == 'on':
-                return xent(score, target) + \
-                       triplet(feat, target)[0] + \
-                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
-                       cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0]  # new add by luo, open label smooth
-            else:
-                return F.cross_entropy(score, target) + \
-                       triplet(feat, target)[0] + \
-                       cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
-                       cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0]  # new add by luo, no label smooth
+                        cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target)
 
         else:
-            print('expected METRIC_LOSS_TYPE with center should be center,'
-                  ' range_center, triplet_center, triplet_range_center '
+            print('expected METRIC_LOSS_TYPE with center should be center, triplet_center'
                   'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
     return loss_func, center_criterion
\ No newline at end of file
diff --git a/layers/cluster_loss.py b/layers/cluster_loss.py
deleted file mode 100644
index 6318389..0000000
--- a/layers/cluster_loss.py
+++ /dev/null
@@ -1,269 +0,0 @@
-from __future__ import absolute_import
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-
-class ClusterLoss(nn.Module):
-    def __init__(self, margin=10, use_gpu=True, ordered=True, ids_per_batch=16, imgs_per_id=4):
-        super(ClusterLoss, self).__init__()
-        self.use_gpu = use_gpu
-        self.margin = margin
-        self.ordered = ordered
-        self.ids_per_batch = ids_per_batch
-        self.imgs_per_id = imgs_per_id
-
-    def _euclidean_dist(self, x, y):
-        """
-        Args:
-          x: pytorch Variable, with shape [m, d]
-          y: pytorch Variable, with shape [n, d]
-        Returns:
-          dist: pytorch Variable, with shape [m, n]
-        """
-        m, n = x.size(0), y.size(0)
-        xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
-        yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
-        dist = xx + yy
-        dist.addmm_(1, -2, x, y.t())
-        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
-        return dist
-
-    def _cluster_loss(self, features, targets, ordered=True, ids_per_batch=16, imgs_per_id=4):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             cluster_loss
-        """
-        if self.use_gpu:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.cpu().unique().cuda()
-            else:
-                unique_labels = targets.cpu().unique().cuda()
-        else:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.unique()
-            else:
-                unique_labels = targets.unique()
-
-        inter_min_distance = torch.zeros(unique_labels.size(0))
-        intra_max_distance = torch.zeros(unique_labels.size(0))
-        center_features = torch.zeros(unique_labels.size(0), features.size(1))
-
-        if self.use_gpu:
-            inter_min_distance = inter_min_distance.cuda()
-            intra_max_distance = intra_max_distance.cuda()
-            center_features = center_features.cuda()
-        
-        index = torch.range(0, unique_labels.size(0) - 1)
-        for i in range(unique_labels.size(0)):
-            label = unique_labels[i]
-            same_class_features = features[targets == label]
-            center_features[i] = same_class_features.mean(dim=0)
-            intra_class_distance = self._euclidean_dist(center_features[index==i], same_class_features)
-            # print('intra_class_distance', intra_class_distance)
-            intra_max_distance[i] = intra_class_distance.max()
-        # print('intra_max_distance:', intra_max_distance)
-
-        for i in range(unique_labels.size(0)):
-            inter_class_distance = self._euclidean_dist(center_features[index==i], center_features[index != i])
-            # print('inter_class_distance', inter_class_distance)
-            inter_min_distance[i] = inter_class_distance.min()
-        #  print('inter_min_distance:', inter_min_distance)
-        cluster_loss = torch.mean(torch.relu(intra_max_distance - inter_min_distance + self.margin))
-        return cluster_loss, intra_max_distance, inter_min_distance
-
-    def forward(self, features, targets):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             cluster_loss
-        """
-        assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
-        cluster_loss, cluster_dist_ap, cluster_dist_an = self._cluster_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
-        return cluster_loss, cluster_dist_ap, cluster_dist_an
-
-
-class ClusterLoss_local(nn.Module):
-    def __init__(self, margin=10, use_gpu=True, ordered=True, ids_per_batch=32, imgs_per_id=4):
-        super(ClusterLoss_local, self).__init__()
-        self.use_gpu = use_gpu
-        self.margin = margin
-        self.ordered = ordered
-        self.ids_per_batch = ids_per_batch
-        self.imgs_per_id = imgs_per_id
-
-    def _euclidean_dist(self, x, y):
-        """
-        Args:
-          x: pytorch Variable, with shape [m, d]
-          y: pytorch Variable, with shape [n, d]
-        Returns:
-          dist: pytorch Variable, with shape [m, n]
-        """
-        m, n = x.size(0), y.size(0)
-        xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
-        yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
-        dist = xx + yy
-        dist.addmm_(1, -2, x, y.t())
-        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
-        return dist
-
-    def _shortest_dist(self, dist_mat):
-        """Parallel version.
-        Args:
-          dist_mat: pytorch Variable, available shape:
-            1) [m, n]
-            2) [m, n, N], N is batch size
-            3) [m, n, *], * can be arbitrary additional dimensions
-        Returns:
-          dist: three cases corresponding to `dist_mat`:
-            1) scalar
-            2) pytorch Variable, with shape [N]
-            3) pytorch Variable, with shape [*]
-        """
-        m, n = dist_mat.size()[:2]
-        # Just offering some reference for accessing intermediate distance.
-        dist = [[0 for _ in range(n)] for _ in range(m)]
-        for i in range(m):
-            for j in range(n):
-                if (i == 0) and (j == 0):
-                    dist[i][j] = dist_mat[i, j]
-                elif (i == 0) and (j > 0):
-                    dist[i][j] = dist[i][j - 1] + dist_mat[i, j]
-                elif (i > 0) and (j == 0):
-                    dist[i][j] = dist[i - 1][j] + dist_mat[i, j]
-                else:
-                    dist[i][j] = torch.min(dist[i - 1][j], dist[i][j - 1]) + dist_mat[i, j]
-        dist = dist[-1][-1]
-        return dist
-
-    def _local_dist(self, x, y):
-        """
-        Args:
-          x: pytorch Variable, with shape [M, m, d]
-          y: pytorch Variable, with shape [N, n, d]
-        Returns:
-          dist: pytorch Variable, with shape [M, N]
-        """
-        M, m, d = x.size()
-        N, n, d = y.size()
-        x = x.contiguous().view(M * m, d)
-        y = y.contiguous().view(N * n, d)
-        # shape [M * m, N * n]
-        dist_mat = self._euclidean_dist(x, y)
-        dist_mat = (torch.exp(dist_mat) - 1.) / (torch.exp(dist_mat) + 1.)
-        # shape [M * m, N * n] -> [M, m, N, n] -> [m, n, M, N]
-        dist_mat = dist_mat.contiguous().view(M, m, N, n).permute(1, 3, 0, 2)
-        # shape [M, N]
-        dist_mat = self._shortest_dist(dist_mat)
-        return dist_mat
-
-    def _cluster_loss(self, features, targets,ordered=True, ids_per_batch=32, imgs_per_id=4):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, H, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             cluster_loss
-        """
-        if self.use_gpu:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.cpu().unique().cuda()
-            else:
-                unique_labels = targets.cpu().unique().cuda()
-        else:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.unique()
-            else:
-                unique_labels = targets.unique()
-
-        inter_min_distance = torch.zeros(unique_labels.size(0))
-        intra_max_distance = torch.zeros(unique_labels.size(0))
-        center_features = torch.zeros(unique_labels.size(0), features.size(1), features.size(2))
-
-        if self.use_gpu:
-            inter_min_distance = inter_min_distance.cuda()
-            intra_max_distance = intra_max_distance.cuda()
-            center_features = center_features.cuda()
-
-        index = torch.range(0, unique_labels.size(0) - 1)
-        for i in range(unique_labels.size(0)):
-            label = unique_labels[i]
-            same_class_features = features[targets == label]
-            center_features[i] = same_class_features.mean(dim=0)
-            intra_class_distance = self._local_dist(center_features[index==i], same_class_features)
-            # print('intra_class_distance', intra_class_distance)
-            intra_max_distance[i] = intra_class_distance.max()
-        # print('intra_max_distance:', intra_max_distance)
-
-        for i in range(unique_labels.size(0)):
-            inter_class_distance = self._local_dist(center_features[index==i], center_features[index != i])
-            # print('inter_class_distance', inter_class_distance)
-            inter_min_distance[i] = inter_class_distance.min()
-        # print('inter_min_distance:', inter_min_distance)
-
-        cluster_loss = torch.mean(torch.relu(intra_max_distance - inter_min_distance + self.margin))
-        return cluster_loss, intra_max_distance, inter_min_distance
-
-    def forward(self, features, targets):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, H, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             cluster_loss
-        """
-        assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
-        cluster_loss, cluster_dist_ap, cluster_dist_an = self._cluster_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
-        return cluster_loss, cluster_dist_ap, cluster_dist_an
-
-
-if __name__ == '__main__':
-    use_gpu = True
-    cluster_loss = ClusterLoss(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
-    features = torch.rand(16, 2048)
-    targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
-    if use_gpu:
-        features = torch.rand(16, 2048).cuda()
-        targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
-    loss = cluster_loss(features, targets)
-    print(loss)
-
-    cluster_loss_local = ClusterLoss_local(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
-    features = torch.rand(16, 8, 2048)
-    targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
-    if use_gpu:
-        features = torch.rand(16, 8, 2048).cuda()
-        targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
-    loss = cluster_loss_local(features, targets)
-    print(loss)
diff --git a/layers/range_loss.py b/layers/range_loss.py
deleted file mode 100644
index de6f587..0000000
--- a/layers/range_loss.py
+++ /dev/null
@@ -1,232 +0,0 @@
-from __future__ import absolute_import
-
-import torch
-from torch import nn
-
-
-class RangeLoss(nn.Module):
-    """
-        Range_loss = alpha * intra_class_loss + beta * inter_class_loss
-        intra_class_loss is the harmonic mean value of the top_k largest distances beturn intra_class_pairs
-        inter_class_loss is the shortest distance between different class centers
-    """
-    def __init__(self, k=2, margin=0.1, alpha=0.5, beta=0.5, use_gpu=True, ordered=True, ids_per_batch=32, imgs_per_id=4):
-        super(RangeLoss, self).__init__()
-        self.use_gpu = use_gpu
-        self.margin = margin
-        self.k = k
-        self.alpha = alpha
-        self.beta = beta
-        self.ordered = ordered
-        self.ids_per_batch = ids_per_batch
-        self.imgs_per_id = imgs_per_id
-
-    def _pairwise_distance(self, features):
-        """
-         Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-         Return: 
-            pairwise distance matrix with shape(batch_size, batch_size)
-        """
-        n = features.size(0)
-        dist = torch.pow(features, 2).sum(dim=1, keepdim=True).expand(n, n)
-        dist = dist + dist.t()
-        dist.addmm_(1, -2, features, features.t())
-        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
-        return dist
-
-    def _compute_top_k(self, features):
-        """
-         Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-         Return: 
-            top_k largest distances
-        """
-        # reading the codes below can help understand better
-        '''
-        dist_array_2 = self._pairwise_distance(features)
-        n = features.size(0)
-        mask = torch.zeros(n, n)
-        if self.use_gpu: mask=mask.cuda()
-        for i in range(0, n):
-            for j in range(i+1, n):
-                mask[i, j] += 1
-        dist_array_2 = dist_array_2 * mask
-        dist_array_2 = dist_array_2.view(1, -1)
-        dist_array_2 = dist_array_2[torch.gt(dist_array_2, 0)]
-        top_k_2 = dist_array_2.sort()[0][-self.k:]
-        print(top_k_2)
-        '''
-        dist_array = self._pairwise_distance(features)
-        dist_array = dist_array.view(1, -1)
-        top_k = dist_array.sort()[0][0, -self.k * 2::2]     # Because there are 2 same value of same feature pair in the dist_array
-        # print('top k intra class dist:', top_k)
-        return top_k
-
-    def _compute_min_dist(self, center_features):
-        """
-         Args:
-            center_features: center matrix (before softmax) with shape (center_number, center_dim)
-         Return: 
-            minimum center distance
-        """
-        '''
-        # reading codes below can help understand better
-        dist_array = self._pairwise_distance(center_features)
-        n = center_features.size(0)
-        mask = torch.zeros(n, n)
-        if self.use_gpu: mask=mask.cuda()
-        for i in range(0, n):
-            for j in range(i + 1, n):
-                mask[i, j] += 1
-        dist_array *= mask
-        dist_array = dist_array.view(1, -1)
-        dist_array = dist_array[torch.gt(dist_array, 0)]
-        min_inter_class_dist = dist_array.min()
-        print(min_inter_class_dist)
-        '''
-        n = center_features.size(0)
-        dist_array2 = self._pairwise_distance(center_features)
-        min_inter_class_dist2 = dist_array2.view(1, -1).sort()[0][0][n]  # exclude self compare, the first one is the min_inter_class_dist
-        return min_inter_class_dist2
-
-    def _calculate_centers(self, features, targets, ordered, ids_per_batch, imgs_per_id):
-        """
-         Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-         Return: 
-            center_features: center matrix (before softmax) with shape (center_number, center_dim)
-        """
-        if self.use_gpu:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.cpu().unique().cuda()
-            else:
-                unique_labels = targets.cpu().unique().cuda()
-        else:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.unique()
-            else:
-                unique_labels = targets.unique()
-
-        center_features = torch.zeros(unique_labels.size(0), features.size(1))
-        if self.use_gpu:
-            center_features = center_features.cuda()
-
-        for i in range(unique_labels.size(0)):
-            label = unique_labels[i]
-            same_class_features = features[targets == label]
-            center_features[i] = same_class_features.mean(dim=0)
-        return center_features
-
-    def _inter_class_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
-        """
-         Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            margin: inter class ringe loss margin
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-         Return: 
-            inter_class_loss
-        """
-        center_features = self._calculate_centers(features, targets, ordered, ids_per_batch, imgs_per_id)
-        min_inter_class_center_distance = self._compute_min_dist(center_features)
-        # print('min_inter_class_center_dist:', min_inter_class_center_distance)
-        return torch.relu(self.margin - min_inter_class_center_distance)
-
-    def _intra_class_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
-        """
-         Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-         Return: 
-            intra_class_loss
-        """
-        if self.use_gpu:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.cpu().unique().cuda()
-            else:
-                unique_labels = targets.cpu().unique().cuda()
-        else:
-            if ordered:
-                if targets.size(0) == ids_per_batch * imgs_per_id:
-                    unique_labels = targets[0:targets.size(0):imgs_per_id]
-                else:
-                    unique_labels = targets.unique()
-            else:
-                unique_labels = targets.unique()
-
-        intra_distance = torch.zeros(unique_labels.size(0))
-        if self.use_gpu:
-            intra_distance = intra_distance.cuda()
-
-        for i in range(unique_labels.size(0)):
-            label = unique_labels[i]
-            same_class_distances = 1.0 / self._compute_top_k(features[targets == label])
-            intra_distance[i] = self.k / torch.sum(same_class_distances)
-        # print('intra_distace:', intra_distance)
-        return torch.sum(intra_distance)
-
-    def _range_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             range_loss
-        """
-        inter_class_loss = self._inter_class_loss(features, targets, ordered, ids_per_batch, imgs_per_id)
-        intra_class_loss = self._intra_class_loss(features, targets, ordered, ids_per_batch, imgs_per_id)
-        range_loss = self.alpha * intra_class_loss + self.beta * inter_class_loss
-        return range_loss, intra_class_loss, inter_class_loss
-
-    def forward(self, features, targets):
-        """
-        Args:
-            features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
-            targets: ground truth labels with shape (batch_size)
-            ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
-            ids_per_batch: num of different ids per batch
-            imgs_per_id: num of images per id
-        Return:
-             range_loss
-        """
-        assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
-        if self.use_gpu:
-            features = features.cuda()
-            targets = targets.cuda()
-
-        range_loss, intra_class_loss, inter_class_loss = self._range_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
-        return range_loss, intra_class_loss, inter_class_loss
-
-
-if __name__ == '__main__':
-    use_gpu = False
-    range_loss = RangeLoss(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
-    features = torch.rand(16, 2048)
-    targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
-    if use_gpu:
-        features = torch.rand(16, 2048).cuda()
-        targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
-    loss = range_loss(features, targets)
-    print(loss)