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