zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

GPU export options (#2297) 

* option for skip last layer and cuda export support

* added parameter device

* fix import

* cleanup 1

* cleanup 2

* opt-in grid

--grid will export with grid computation, default export will skip grid (same as current)

* default --device cpu

GPU export causes ONNX and CoreML errors.

Co-authored-by: Jan Hajek <jan.hajek@gmail.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Jan Hajek 2021-03-06 21:02:10 +01:00 committed by GitHub
parent cd8ed3521d
commit 7a0a81fd1d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
models/export.py
View File

@ -17,13 +17,16 @@ import models
from models.experimental import attempt_load from models.experimental import attempt_load
from utils.activations import Hardswish, SiLU from utils.activations import Hardswish, SiLU
from utils.general import set_logging, check_img_size from utils.general import set_logging, check_img_size
from utils.torch_utils import select_device
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/
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('--dynamic', action='store_true', help='dynamic ONNX axes')
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('--dynamic', action='store_true', help='dynamic ONNX axes')
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')
opt = parser.parse_args() opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
print(opt) print(opt)
@ -31,7 +34,8 @@ if __name__ == '__main__':
t = time.time() t = time.time()
# Load PyTorch model # Load PyTorch model
model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model device = select_device(opt.device)
model = attempt_load(opt.weights, map_location=device) # load FP32 model
labels = model.names labels = model.names
# Checks # Checks
@ -39,7 +43,7 @@ if __name__ == '__main__':
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
# Input # Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size) # image size(1,3,320,192) iDetection img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
# Update model # Update model
for k, m in model.named_modules(): for k, m in model.named_modules():
@ -51,7 +55,7 @@ if __name__ == '__main__':
m.act = SiLU() m.act = SiLU()
# elif isinstance(m, models.yolo.Detect): # elif isinstance(m, models.yolo.Detect):
# m.forward = m.forward_export # assign forward (optional) # m.forward = m.forward_export # assign forward (optional)
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = not opt.grid # set Detect() layer grid export
y = model(img) # dry run y = model(img) # dry run
# TorchScript export # TorchScript export