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

ONNX compatibility updates

This commit is contained in:
Glenn Jocher 2018-12-28 20:15:26 +01:00
parent 16bc3b72c3
commit d1951c1868
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
models.py
View File

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