From efbeb283c4953501f31d005619ced9bef7886a2c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 26 Apr 2020 14:01:20 -0700 Subject: [PATCH] ONNX grid float --- models.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models.py b/models.py index 9d0d7de5..5e1ac1b3 100755 --- a/models.py +++ b/models.py @@ -145,7 +145,7 @@ class YOLOLayer(nn.Module): def create_grids(self, ng=(13, 13), device='cpu'): self.nx, self.ny = ng # x and y grid size - self.ng = torch.tensor(ng) + self.ng = torch.tensor(ng, dtype=torch.float) # build xy offsets if not self.training: @@ -193,9 +193,9 @@ class YOLOLayer(nn.Module): elif ONNX_EXPORT: # Avoid broadcasting for ANE operations m = self.na * self.nx * self.ny - ng = 1 / self.ng.repeat((m, 1)) - grid = self.grid.repeat((1, self.na, 1, 1, 1)).view(m, 2) - anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(m, 2) * ng + ng = 1. / self.ng.repeat(m, 1) + grid = self.grid.repeat(1, self.na, 1, 1, 1).view(m, 2) + anchor_wh = self.anchor_wh.repeat(1, 1, self.nx, self.ny, 1).view(m, 2) * ng p = p.view(m, self.no) xy = torch.sigmoid(p[:, 0:2]) + grid # x, y