yjh0410 1 gadu atpakaļ
vecāks
revīzija
3da43701c2
1 mainītis faili ar 3 papildinājumiem un 0 dzēšanām
  1. 3 0
      yolo/models/yolov5/yolov5_pred.py

+ 3 - 0
yolo/models/yolov5/yolov5_pred.py

@@ -45,6 +45,9 @@ class DetPredLayer(nn.Module):
         b = self.reg_pred.bias.view(-1, )
         b.data.fill_(1.0)
         self.reg_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+        w = self.reg_pred.weight
+        w.data.fill_(0.)
+        self.reg_pred.weight = torch.nn.Parameter(w, requires_grad=True)
 
     def generate_anchors(self, fmp_size):
         """