Skip to content

Commit

Permalink
Add grid concat and fuse such operators (WongKinYiu#389)
Browse files Browse the repository at this point in the history 
* Add grid concat and fuse so many op

* Fix model

* Fix other detector

* Update yolo.py

* Update yolo.py

Co-authored-by: Alexey <AlexeyAB@users.noreply.github.com>
  • Loading branch information
@triple-Mu @AlexeyAB
triple-Mu and AlexeyAB authored Aug 9, 2022
1 parent c14ba0c commit 8dc755a
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 19 deletions.
19 changes: 11 additions & 8 deletions export.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,15 +143,18 @@
'output': {0: 'batch'},
}
dynamic_axes.update(output_axes)
if opt.grid and opt.end2end:
print('\nStarting export end2end onnx model for %s...' % 'TensorRT' if opt.max_wh is None else 'onnxruntime')
model = End2End(model,opt.topk_all,opt.iou_thres,opt.conf_thres,opt.max_wh,device)
if opt.end2end and opt.max_wh is None:
output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']
shapes = [opt.batch_size, 1, opt.batch_size, opt.topk_all, 4,
opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
if opt.grid:
if opt.end2end:
print('\nStarting export end2end onnx model for %s...' % 'TensorRT' if opt.max_wh is None else 'onnxruntime')
model = End2End(model,opt.topk_all,opt.iou_thres,opt.conf_thres,opt.max_wh,device)
if opt.end2end and opt.max_wh is None:
output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']
shapes = [opt.batch_size, 1, opt.batch_size, opt.topk_all, 4,
opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
else:
output_names = ['output']
else:
output_names = ['output']
model.model[-1].concat = True

torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
output_names=output_names,
Expand Down
34 changes: 23 additions & 11 deletions models/yolo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ class Detect(nn.Module):
stride = None # strides computed during build
export = False # onnx export
end2end = False
include_nms = False
include_nms = False
concat = False

def __init__(self, nc=80, anchors=(), ch=()): # detection layer
super(Detect, self).__init__()
Expand Down Expand Up @@ -55,9 +56,10 @@ def forward(self, x):
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else:
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
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
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, -1, self.no))

if self.training:
Expand All @@ -67,6 +69,8 @@ def forward(self, x):
elif self.include_nms:
z = self.convert(z)
out = (z, )
elif self.concat:
out = torch.cat(z, 1)
else:
out = (torch.cat(z, 1), x)

Expand Down Expand Up @@ -94,7 +98,8 @@ class IDetect(nn.Module):
stride = None # strides computed during build
export = False # onnx export
end2end = False
include_nms = False
include_nms = False
concat = False

def __init__(self, nc=80, anchors=(), ch=()): # detection layer
super(IDetect, self).__init__()
Expand Down Expand Up @@ -150,9 +155,10 @@ def fuseforward(self, x):
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else:
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
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
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, -1, self.no))

if self.training:
Expand All @@ -162,6 +168,8 @@ def fuseforward(self, x):
elif self.include_nms:
z = self.convert(z)
out = (z, )
elif self.concat:
out = torch.cat(z, 1)
else:
out = (torch.cat(z, 1), x)

Expand Down Expand Up @@ -305,6 +313,7 @@ class IAuxDetect(nn.Module):
export = False # onnx export
end2end = False
include_nms = False
concat = False

def __init__(self, nc=80, anchors=(), ch=()): # detection layer
super(IAuxDetect, self).__init__()
Expand Down Expand Up @@ -344,9 +353,10 @@ def forward(self, x):
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else:
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
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
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, -1, self.no))

return x if self.training else (torch.cat(z, 1), x[:self.nl])
Expand Down Expand Up @@ -381,6 +391,8 @@ def fuseforward(self, x):
elif self.include_nms:
z = self.convert(z)
out = (z, )
elif self.concat:
out = torch.cat(z, 1)
else:
out = (torch.cat(z, 1), x)

Expand Down

0 comments on commit 8dc755a

Please sign in to comment.