From d1951c1868ce0cd39821911a387cb6e805b7174a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 28 Dec 2018 20:15:26 +0100 Subject: [PATCH] ONNX compatibility updates --- models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/models.py b/models.py index a08cd822..565cecda 100755 --- a/models.py +++ b/models.py @@ -317,6 +317,7 @@ class Darknet(nn.Module): ONNX_export = False if ONNX_export: + # Produce a single-layer *.onnx model (upsample ops not working in PyTorch 1.0 export yet) output = output[0].squeeze().transpose(0, 1) # first layer reshaped to 85 x 507 output[5:] = torch.nn.functional.softmax(torch.sigmoid(output[5:]) * output[4:5], dim=0) # SSD-like conf return output[5:], output[:4] # ONNX scores, boxes