Add support for TFLite INT8 detection export

README.md

@@ -183,6 +183,15 @@ python ./tensorrt-python/export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p
 
 Tested with: Python 3.7.13, Pytorch 1.12.0+cu113
 
+**Pytorch to TFLite INT8** <a href="https://colab.research.google.com/github/WongKinYiu/yolov7/blob/main/tools/YOLOv7tflite_int8.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
+```shell
+python export.py --weights yolov7-tiny.pt --grid --simplify --img-size 640 640 --normalize
+onnx2tf -i yolov7-tiny.onnx -o yolov7-tiny.tf --verbosity info -nuo -oiqt -qt per-tensor
+# Or use docker to export with onnx2tf
+docker run --rm -v `pwd`:/workdir -w /workdir docker.io/pinto0309/onnx2tf:1.19.4 \
+    onnx2tf -i yolov7-tiny.onnx -o yolov7-tiny.tf --verbosity info -nuo -oiqt -qt per-tensor
+```
+
 ## Pose estimation
 
 [`code`](https://github.com/WongKinYiu/yolov7/tree/pose) [`yolov7-w6-pose.pt`](https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-w6-pose.pt)

export.py

@@ -34,6 +34,7 @@
     parser.add_argument('--include-nms', action='store_true', help='export end2end onnx')
     parser.add_argument('--fp16', action='store_true', help='CoreML FP16 half-precision export')
     parser.add_argument('--int8', action='store_true', help='CoreML INT8 quantization')
+    parser.add_argument('--normalize', action='store_true', help='Normal pixel coordinate output of head')
     opt = parser.parse_args()
     opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
     opt.dynamic = opt.dynamic and not opt.end2end
@@ -69,6 +70,9 @@
     if opt.include_nms:
         model.model[-1].include_nms = True
         y = None
+    if opt.normalize:
+        model.model[-1].normalize = True
+        model.model[-1].imgszxy = torch.tensor([opt.img_size[1], opt.img_size[0]])
 
     # TorchScript export
     try:

models/yolo.py

@@ -26,6 +26,8 @@ class Detect(nn.Module):
     end2end = False
     include_nms = False
     concat = False
+    normalize = False  # set only for export, normalize pixel value outpus
+    imgszxy = None  # set only for export, static model input image size
 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(Detect, self).__init__()
@@ -57,8 +59,15 @@ def forward(self, x):
                     y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                 else:
                     xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
-                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    if not self.normalize:
+                        xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
+                        wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    else:
+                        # normalized output of pixel coordinates
+                        # precompute normalized constants for static models
+                        # multiplication will be more stable for static quantized models
+                        xy = xy / (self.imgszxy / (2. * self.stride[i])) + (self.stride[i] * (self.grid[i] - 0.5)) / self.imgszxy  # new xy
+                        wh = wh ** 2 * ((4 * self.anchor_grid[i].data) / self.imgszxy)  # new wh
                     y = torch.cat((xy, wh, conf), 4)
                 z.append(y.view(bs, -1, self.no))
 
@@ -100,6 +109,8 @@ class IDetect(nn.Module):
     end2end = False
     include_nms = False
     concat = False
+    normalize = False  # set only for export, normalize pixel value outpus
+    imgszxy = None  # set only for export, static model input image size
 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(IDetect, self).__init__()
@@ -156,8 +167,15 @@ def fuseforward(self, x):
                     y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                 else:
                     xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
-                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    if not self.normalize:
+                        xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
+                        wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    else:
+                        # normalized output of pixel coordinates
+                        # precompute normalized constants for static models
+                        # multiplication will be more stable for static quantized models
+                        xy = xy / (self.imgszxy / (2. * self.stride[i])) + (self.stride[i] * (self.grid[i] - 0.5)) / self.imgszxy  # new xy
+                        wh = wh ** 2 * ((4 * self.anchor_grid[i].data) / self.imgszxy)  # new wh
                     y = torch.cat((xy, wh, conf), 4)
                 z.append(y.view(bs, -1, self.no))