From 841fa1452440349220187d58f41a893bb1e890ee Mon Sep 17 00:00:00 2001 From: wa22 Date: Thu, 23 May 2024 13:11:50 +0000 Subject: [PATCH] update --- README.md | 2 +- ultralytics/nn/modules/head.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ccfba4b80..82c4f1485 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,7 @@ yolo predict model=yolov10n/s/m/b/l/x.engine ## Acknowledgement -The code base is built with [ultralytics](https://github.com/ultralytics/ultralytics) and [RT-DETR](https://github.com/lyuwenyu/RT-DETR) +The code base is built with [ultralytics](https://github.com/ultralytics/ultralytics) and [RT-DETR](https://github.com/lyuwenyu/RT-DETR). Thanks for the great implementations! diff --git a/ultralytics/nn/modules/head.py b/ultralytics/nn/modules/head.py index 4d6816e34..ba9d2d218 100644 --- a/ultralytics/nn/modules/head.py +++ b/ultralytics/nn/modules/head.py @@ -510,6 +510,9 @@ def __init__(self, nc=80, ch=()): def forward(self, x): one2one = self.forward_feat([xi.detach() for xi in x], self.one2one_cv2, self.one2one_cv3) + if not self.export: + one2many = super().forward(x) + if not self.training: one2one = self.inference(one2one) if not self.export: @@ -519,7 +522,6 @@ def forward(self, x): boxes, scores, labels = ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det) return torch.cat([boxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1) else: - one2many = super().forward(x) return {"one2many": one2many, "one2one": one2one} def bias_init(self):