From 868c0e9bbb45b031e7bfd73c6d3983bcce07b9c1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 19 Sep 2022 13:31:24 +0200 Subject: [PATCH] Update Detect() grid init `for` loop (#9494) May resolve threaded inference issue in https://github.com/ultralytics/yolov5/pull/9425#issuecomment-1250802928 by avoiding memory sharing on init. Signed-off-by: Glenn Jocher Signed-off-by: Glenn Jocher --- models/yolo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 0dca635..1d0da2a 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -47,8 +47,8 @@ class Detect(nn.Module): self.no = nc + 5 # number of outputs per anchor self.nl = len(anchors) # number of detection layers self.na = len(anchors[0]) // 2 # number of anchors - self.grid = [torch.empty(1)] * self.nl # init grid - self.anchor_grid = [torch.empty(1)] * self.nl # init anchor grid + self.grid = [torch.empty(0) for _ in range(self.nl)] # init grid + self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] # init anchor grid self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2) self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv self.inplace = inplace # use inplace ops (e.g. slice assignment)