Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

ONNX grid float

This commit is contained in:
Glenn Jocher 2020-04-26 14:01:20 -07:00
parent 3bf0cb9c60
commit efbeb283c4
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
models.py
View File

@ -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