From 003de1917ec6135b70ad08a1ecba0a1c6ea7ca8b Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Wed, 13 Mar 2019 12:50:13 +1100 Subject: [PATCH 1/2] Fix shape-mismatch in ONNX export --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models.py b/models.py index a2db26e9..3d3acad3 100755 --- a/models.py +++ b/models.py @@ -190,7 +190,7 @@ class YOLOLayer(nn.Module): # return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t() 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 p_conf = torch.sigmoid(p[..., 4:5]) # Conf p_cls = p[..., 5:85] From 10182ca39d1ebc6b1f36994c7c495ae9d38e3fb6 Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Wed, 13 Mar 2019 12:51:05 +1100 Subject: [PATCH 2/2] Get rid of hardcoded values of 85 --- models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models.py b/models.py index 3d3acad3..c72c10a4 100755 --- a/models.py +++ b/models.py @@ -189,11 +189,11 @@ class YOLOLayer(nn.Module): # p_cls = F.softmax(p[:, 5:85], 1) * p_conf # SSD-like conf # return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t() - p = p.view(1, -1, 85) + p = p.view(1, -1, 5 + self.nC) xy = xy.view(bs, self.nA * nG * nG, 2) + grid_xy # x, y wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height p_conf = torch.sigmoid(p[..., 4:5]) # Conf - p_cls = p[..., 5:85] + p_cls = p[..., 5:5 + self.nC] # Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py # p_cls = F.softmax(p_cls, 2) * p_conf # SSD-like conf p_cls = torch.exp(p_cls).permute((2, 1, 0)) @@ -260,7 +260,7 @@ class Darknet(nn.Module): if ONNX_EXPORT: output = torch.cat(output, 1) # merge the 3 layers 85 x (507, 2028, 8112) to 85 x 10647 - return output[5:85].t(), output[:4].t() # ONNX scores, boxes + return output[5:].t(), output[:4].t() # ONNX scores, boxes return sum(output) if is_training else torch.cat(output, 1)