Merge branch 'master' of https://github.com/Tianxiaomo/pytorch-YOLOv4

Author: wenh06

.gitignore

@@ -5,6 +5,7 @@ ttest
 *.engine
 *.pyc
 *.infer
+*.npy
 
 z_demo_*
 

README.md

@@ -119,7 +119,7 @@ you can use darknet2pytorch to convert it yourself, or download my converted mod
 
 - **This script is to convert the official pretrained darknet model into ONNX**
 
-- **Pytorch version Recommended: 1.4.0**
+- **Pytorch version Recommended: 1.4.0 for TensorRT 7.0, 1.5.0 or higher for TensorRT 7.1**
 
 - **Install onnxruntime**
 
@@ -142,7 +142,7 @@ you can use darknet2pytorch to convert it yourself, or download my converted mod
 
 - **You can convert your trained pytorch model into ONNX using this script**
 
-- **Pytorch version Recommended: 1.4.0**
+- **Pytorch version Recommended: 1.4.0 for TensorRT 7.0, 1.5.0 or higher for TensorRT 7.1**
 
 - **Install onnxruntime**
 
@@ -191,8 +191,6 @@ you can use darknet2pytorch to convert it yourself, or download my converted mod
     
     - Note2: extra NMS operations are needed for the tensorRT output. This demo uses python NMS code from `tool/utils.py`.
 
-    - Inference is verified to be okay for TensorRT 7.0, but TensorRT 7.1 may be problematic in handling y coordinates of bounding boxes. Problem investigation is underway.
-
 
 # 6. ONNX2Tensorflow
 

demo_trt.py

@@ -119,7 +119,7 @@ def main(engine_path, image_path, image_size):
 
         for i in range(2):  # This 'for' loop is for speed check
                             # Because the first iteration is usually longer
-            boxes = detect(engine, context, buffers, image_src, image_size, num_classes)
+            boxes = detect(context, buffers, image_src, image_size, num_classes)
 
         if num_classes == 20:
             namesfile = 'data/voc.names'
@@ -140,7 +140,7 @@ def get_engine(engine_path):
 
 
 
-def detect(engine, context, buffers, image_src, image_size, num_classes):
+def detect(context, buffers, image_src, image_size, num_classes):
     IN_IMAGE_H, IN_IMAGE_W = image_size
 
     ta = time.time()

tool/yolo_layer.py

@@ -189,13 +189,13 @@ def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x
     for i in range(num_anchors):
         ii = i * 2
         # Shape: [batch, 1, H, W]
-        bx = bxy[:, ii] + torch.tensor(grid_x, device=device, dtype=torch.float32) # grid_x.to(device=device, dtype=torch.float32)
+        bx = bxy[:, ii : ii + 1] + torch.tensor(grid_x, device=device, dtype=torch.float32) # grid_x.to(device=device, dtype=torch.float32)
         # Shape: [batch, 1, H, W]
-        by = bxy[:, ii + 1] + torch.tensor(grid_y, device=device, dtype=torch.float32) # grid_y.to(device=device, dtype=torch.float32)
+        by = bxy[:, ii + 1 : ii + 2] + torch.tensor(grid_y, device=device, dtype=torch.float32) # grid_y.to(device=device, dtype=torch.float32)
         # Shape: [batch, 1, H, W]
-        bw = bwh[:, ii] * anchor_w[i]
+        bw = bwh[:, ii : ii + 1] * anchor_w[i]
         # Shape: [batch, 1, H, W]
-        bh = bwh[:, ii + 1] * anchor_h[i]
+        bh = bwh[:, ii + 1 : ii + 2] * anchor_h[i]
 
         bx_list.append(bx)
         by_list.append(by)
@@ -216,17 +216,20 @@ def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x
     # Shape: [batch, num_anchors, H, W]
     bh = torch.cat(bh_list, dim=1)
 
+    # Shape: [batch, 2 * num_anchors, H, W]
+    bx_bw = torch.cat((bx, bw), dim=1)
+    # Shape: [batch, 2 * num_anchors, H, W]
+    by_bh = torch.cat((by, bh), dim=1)
+
     # normalize coordinates to [0, 1]
-    bx = bx / W
-    by = by / H
-    bw = bw / W
-    bh = bh / H
+    bx_bw /= W
+    by_bh /= H
 
     # Shape: [batch, num_anchors * H * W, 1]
-    bx = bx.view(batch, num_anchors * H * W, 1)
-    by = by.view(batch, num_anchors * H * W, 1)
-    bw = bw.view(batch, num_anchors * H * W, 1)
-    bh = bh.view(batch, num_anchors * H * W, 1)
+    bx = bx_bw[:, :num_anchors].view(batch, num_anchors * H * W, 1)
+    by = by_bh[:, :num_anchors].view(batch, num_anchors * H * W, 1)
+    bw = bx_bw[:, num_anchors:].view(batch, num_anchors * H * W, 1)
+    bh = by_bh[:, num_anchors:].view(batch, num_anchors * H * W, 1)
 
     # Shape: [batch, num_anchors * h * w, 4]
     boxes = torch.cat((bx, by, bw, bh), dim=2).view(batch, num_anchors * H * W, 4)