mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Fix ONNX export using --grid --simplify --dynamic simultaneously (#2982)

Browse Source 
* Update yolo.py

* Update export.py

* fix export grid

* Update export.py, remove detect export attribute

* rearrange if order

* remove --grid, default inplace=False

* rename exp_dynamic to onnx_dynamic, comment

* replace bs with 1 in anchor_grid[i] index 0

* Update export.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
jylink 2021-05-04 01:01:29 +08:00 committed by GitHub
parent 41cc7caee6
commit b292837e36
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
models/export.py
View File

@ -26,9 +26,9 @@ if __name__ == '__main__':
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', action='store_true', help='FP16 half-precision export') parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
opt = parser.parse_args() opt = parser.parse_args()
@ -60,9 +60,11 @@ if __name__ == '__main__':
m.act = Hardswish() m.act = Hardswish()
elif isinstance(m.act, nn.SiLU): elif isinstance(m.act, nn.SiLU):
m.act = SiLU() m.act = SiLU()
# elif isinstance(m, models.yolo.Detect): elif isinstance(m, models.yolo.Detect):
m.inplace = opt.inplace
m.onnx_dynamic = opt.dynamic
# m.forward = m.forward_export # assign forward (optional) # m.forward = m.forward_export # assign forward (optional)
model.model[-1].export = not opt.grid # set Detect() layer grid export
for _ in range(2): for _ in range(2):
y = model(img) # dry runs y = model(img) # dry runs
print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)") print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")

7
models/yolo.py
View File

@ -24,7 +24,7 @@ except ImportError:
class Detect(nn.Module): class Detect(nn.Module):
stride = None # strides computed during build stride = None # strides computed during build
export = False # onnx export onnx_dynamic = False # ONNX export parameter
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
super(Detect, self).__init__() super(Detect, self).__init__()
@ -42,14 +42,13 @@ class Detect(nn.Module):
def forward(self, x): def forward(self, x):
# x = x.copy() # for profiling # x = x.copy() # for profiling
z = [] # inference output z = [] # inference output
self.training |= self.export
for i in range(self.nl): for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4]: if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
self.grid[i] = self._make_grid(nx, ny).to(x[i].device) self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
y = x[i].sigmoid() y = x[i].sigmoid()
@ -58,7 +57,7 @@ class Detect(nn.Module):
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953 else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
y = torch.cat((xy, wh, y[..., 4:]), -1) y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no)) z.append(y.view(bs, -1, self.no))