Upload code&model for NHWC format

onnx_eval.py

@@ -3,9 +3,10 @@ import json
 import os
 import sys
 from pathlib import Path
+import torch
 import onnxruntime
 import numpy as np
-import torch
+
 from tqdm import tqdm
 from pycocotools.coco import COCO
 from pycocotools.cocoeval import COCOeval
@@ -145,8 +146,10 @@ def run(data,
         targets = targets.to(device)
         nb, _, height, width = img.shape  # batch size, channels, height, width
 
-        outputs = onnx_model.run(None, {onnx_model.get_inputs()[0].name: img.cpu().numpy()})
-        outputs = [torch.tensor(item).to(device) for item in outputs]
+        # outputs = onnx_model.run(None, {onnx_model.get_inputs()[0].name: img.cpu().numpy()})
+        outputs = onnx_model.run(None, {onnx_model.get_inputs()[0].name: img.permute(0, 2, 3, 1).cpu().numpy()})
+        # outputs = [torch.tensor(item).to(device) for item in outputs]
+        outputs = [torch.tensor(item).permute(0, 3, 1, 2).to(device) for item in outputs]
         outputs = post_process(outputs)
         out, train_out = outputs[0], outputs[1]
 
@@ -267,4 +270,4 @@ def main(opt):
 
 if __name__ == "__main__":
     opt = parse_opt()
-    main(opt)
+    main(opt)

onnx_inference.py

@@ -1,7 +1,8 @@
-import onnxruntime
+
 import numpy as np
 import cv2
 import torch
+import onnxruntime
 import sys
 import pathlib
 CURRENT_DIR = pathlib.Path(__file__).parent
@@ -113,8 +114,9 @@ if __name__ == '__main__':
 
     img0 = cv2.imread(path)
     img = pre_process(img0)
-    onnx_input = {onnx_model.get_inputs()[0].name: img}
+    onnx_input = {onnx_model.get_inputs()[0].name: img.transpose(0, 2, 3, 1)}
     onnx_output = onnx_model.run(None, onnx_input)
+    onnx_output = [torch.tensor(item).permute(0, 3, 1, 2) for item in onnx_output]
     onnx_output = post_process(onnx_output)
     pred = non_max_suppression(
         onnx_output[0], conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det

yolov5s_qat.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5ba00d5f170eab6130610bb543c1f4b1e8354f4944c127e61c28beb99beddf26
-size 29141657
+oid sha256:5f05e2860614a4d10757405f5e4ad2849d380631e16915f91aa0f69597d10575
+size 29142007