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

Fix shape-mismatch in ONNX export

This commit is contained in:
Daniel Suess 2019-03-13 12:50:13 +11:00
parent c1c09eb3cc
commit 003de1917e
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
models.py
View File

@ -190,7 +190,7 @@ class YOLOLayer(nn.Module):
# return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t() # return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t()
p = p.view(1, -1, 85) p = p.view(1, -1, 85)
xy = xy + grid_xy # x, y xy = xy.view(bs, self.nA * nG * nG, 2) + grid_xy # x, y
wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height
p_conf = torch.sigmoid(p[..., 4:5]) # Conf p_conf = torch.sigmoid(p[..., 4:5]) # Conf
p_cls = p[..., 5:85] p_cls = p[..., 5:85]