fix ReidEvalution too much memory cost

Move `matches` matrix computation in each iteration to reduce the extra memory cost #420 #404
JDAI-CV · Jun 8, 2021 · 7ed6240 · 7ed6240
1 parent de81b3d
commit 7ed6240
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 deletions.
diff --git a/fastreid/evaluation/rank.py b/fastreid/evaluation/rank.py
@@ -107,9 +107,6 @@ def eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):
         print('Note: number of gallery samples is quite small, got {}'.format(num_g))
 
     indices = np.argsort(distmat, axis=1)
-
-    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)
-
     # compute cmc curve for each query
     all_cmc = []
     all_AP = []
@@ -127,7 +124,8 @@ def eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):
         keep = np.invert(remove)
 
         # compute cmc curve
-        raw_cmc = matches[q_idx][keep]  # binary vector, positions with value 1 are correct matches
+        matches = (g_pids[order] == q_pid).astype(np.int32)
+        raw_cmc = matches[keep]  # binary vector, positions with value 1 are correct matches
         if not np.any(raw_cmc):
             # this condition is true when query identity does not appear in gallery
             continue
@@ -163,7 +161,7 @@ def eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):
 
 def evaluate_py(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03):
     if use_metric_cuhk03:
-        return eval_cuhk03(distmat, g_pids, q_camids, g_camids, max_rank)
+        return eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)
     else:
         return eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)
 
@@ -176,7 +174,7 @@ def evaluate_rank(
         g_camids,
         max_rank=50,
         use_metric_cuhk03=False,
-        use_cython=True
+        use_cython=True,
 ):
     """Evaluates CMC rank.
     Args:

diff --git a/fastreid/evaluation/rank_cylib/rank_cy.pyx b/fastreid/evaluation/rank_cylib/rank_cy.pyx
@@ -159,7 +159,7 @@ cpdef eval_market1501_cy(float[:,:] distmat, long[:] q_pids, long[:]g_pids,
 
     cdef:
         long[:,:] indices = np.argsort(distmat, axis=1)
-        long[:,:] matches = (np.asarray(g_pids)[np.asarray(indices)] == np.asarray(q_pids)[:, np.newaxis]).astype(np.int64)
+        long[:] matches
 
         float[:,:] all_cmc = np.zeros((num_q, max_rank), dtype=np.float32)
         float[:] all_AP = np.zeros(num_q, dtype=np.float32)
@@ -192,14 +192,15 @@ cpdef eval_market1501_cy(float[:,:] distmat, long[:] q_pids, long[:]g_pids,
             order[g_idx] = indices[q_idx, g_idx]
         num_g_real = 0
         meet_condition = 0
+        matches = (np.asarray(g_pids)[np.asarray(order)] == q_pid).astype(np.int64)
 
         # remove gallery samples that have the same pid and camid with query
         for g_idx in range(num_g):
             if (g_pids[order[g_idx]] != q_pid) or (g_camids[order[g_idx]] != q_camid):
-                raw_cmc[num_g_real] = matches[q_idx][g_idx]
+                raw_cmc[num_g_real] = matches[g_idx]
                 num_g_real += 1
                 # this condition is true if query appear in gallery
-                if matches[q_idx][g_idx] > 1e-31:
+                if matches[g_idx] > 1e-31:
                     meet_condition = 1
 
         if not meet_condition:

diff --git a/fastreid/layers/batch_norm.py b/fastreid/layers/batch_norm.py
@@ -185,7 +185,7 @@ def get_norm(norm, out_channels, **kwargs):
     """
     Args:
         norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;
-            or a callable that thakes a channel number and returns
+            or a callable that takes a channel number and returns
             the normalization layer as a nn.Module
         out_channels: number of channels for normalization layer