From fa6c5a11820e550ce31419d83960b8fcd18fc748 Mon Sep 17 00:00:00 2001
From: Guanghua Yu <742925032@qq.com>
Date: Fri, 28 May 2021 12:19:28 +0800
Subject: [PATCH] fix num_classes in solov2 (#3193)

---
 ppdet/modeling/heads/solov2_head.py               | 3 ++-
 static/ppdet/modeling/anchor_heads/solov2_head.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/ppdet/modeling/heads/solov2_head.py b/ppdet/modeling/heads/solov2_head.py
index 9104b7d2a..ac5ebe7df 100644
--- a/ppdet/modeling/heads/solov2_head.py
+++ b/ppdet/modeling/heads/solov2_head.py
@@ -462,7 +462,8 @@ class SOLOv2Head(nn.Layer):
         # cate_labels & kernel_preds
         cate_labels = inds[:, 1]
         kernel_preds = paddle.gather(kernel_preds, index=inds[:, 0])
-        cate_score_idx = paddle.add(inds[:, 0] * 80, cate_labels)
+        cate_score_idx = paddle.add(inds[:, 0] * self.cate_out_channels,
+                                    cate_labels)
         cate_scores = paddle.gather(cate_preds, index=cate_score_idx)
 
         size_trans = np.power(self.seg_num_grids, 2)
diff --git a/static/ppdet/modeling/anchor_heads/solov2_head.py b/static/ppdet/modeling/anchor_heads/solov2_head.py
index 17473110d..590b428e8 100644
--- a/static/ppdet/modeling/anchor_heads/solov2_head.py
+++ b/static/ppdet/modeling/anchor_heads/solov2_head.py
@@ -367,8 +367,8 @@ class SOLOv2Head(object):
         # cate_labels & kernel_preds
         cate_labels = inds[:, 1]
         kernel_preds = fluid.layers.gather(kernel_preds, index=inds[:, 0])
-        cate_score_idx = fluid.layers.elementwise_add(inds[:, 0] * 80,
-                                                      cate_labels)
+        cate_score_idx = fluid.layers.elementwise_add(
+            inds[:, 0] * self.cate_out_channels, cate_labels)
         cate_scores = fluid.layers.gather(cate_preds, index=cate_score_idx)
 
         size_trans = np.power(self.seg_num_grids, 2)
-- 
GitLab