From c0a126bf84a6ca3a13ba796f76d4e8721d68312f Mon Sep 17 00:00:00 2001
From: whitecrow5 <35173050+WhiteCrow5@users.noreply.github.com>
Date: Fri, 15 Oct 2021 14:45:14 +0800
Subject: [PATCH] Update rr_yolov5_head.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

修正类别非80时产生的bug
---
 mmdet/models/dense_heads/rr_yolov5_head.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mmdet/models/dense_heads/rr_yolov5_head.py b/mmdet/models/dense_heads/rr_yolov5_head.py
index b27f3bed..44a94464 100644
--- a/mmdet/models/dense_heads/rr_yolov5_head.py
+++ b/mmdet/models/dense_heads/rr_yolov5_head.py
@@ -70,9 +70,9 @@ def _init_layers(self):
 
         self.det = nn.Sequential(*model)
         self.head = nn.Sequential(
-            nn.Conv2d(make_div8_fun(256), 255, 1),
-            nn.Conv2d(make_div8_fun(512), 255, 1),
-            nn.Conv2d(make_div8_fun(1024), 255, 1),
+            nn.Conv2d(make_div8_fun(256), (self.num_classes + 5) * 3, 1),
+            nn.Conv2d(make_div8_fun(512), (self.num_classes + 5) * 3, 1),
+            nn.Conv2d(make_div8_fun(1024), (self.num_classes + 5) * 3, 1),
         )
 
     def forward(self, feats):