瀏覽代碼

optimize pred layer

yjh0410 2 年之前
父節點
當前提交
6314e306f6
共有 1 個文件被更改,包括 1 次插入0 次删除
  1. 1 0
      models/detectors/yolov8/yolov8_pred.py

+ 1 - 0
models/detectors/yolov8/yolov8_pred.py

@@ -141,6 +141,7 @@ class MultiLevelPredLayer(nn.Module):
             x2y2_pred = outputs["anchors"][None] + delta_pred[..., 2:] * self.strides[level]
             box_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
 
+            # collect results
             all_cls_preds.append(outputs["pred_cls"])
             all_reg_preds.append(outputs["pred_reg"])
             all_box_preds.append(box_pred)