mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Earlier `assert` for cpu and half option (#3508) 

* early assert for cpu and half option

early assert for cpu and half option

* Modified comment

Modified comment
pull/3510/head
Yonghye Kwon 2021-06-08 01:56:41 +09:00 committed by GitHub
parent 90b7895d65
commit 8d1ddc93c7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
models/export.py
View File

@ -44,15 +44,13 @@ if __name__ == '__main__':
# Load PyTorch model # Load PyTorch model
device = select_device(opt.device) device = select_device(opt.device)
assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
model = attempt_load(opt.weights, map_location=device) # load FP32 model model = attempt_load(opt.weights, map_location=device) # load FP32 model
labels = model.names labels = model.names
# Checks # Input
gs = int(max(model.stride)) # grid size (max stride) gs = int(max(model.stride)) # grid size (max stride)
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
assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
# Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # 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