(
data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
weights=ROOT / 'yolov5s.pt', # weights path
imgsz=(640, 640), # image (height, width)
batch_size=1, # batch size
device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
include=('torchscript', 'onnx'), # include formats
half=False, # FP16 half-precision export
inplace=False, # set YOLOv5 Detect() inplace=True
keras=False, # use Keras
optimize=False, # TorchScript: optimize for mobile
int8=False, # CoreML/TF INT8 quantization
dynamic=False, # ONNX/TF/TensorRT: dynamic axes
simplify=False, # ONNX: simplify model
opset=12, # ONNX: opset version
verbose=False, # TensorRT: verbose log
workspace=4, # TensorRT: workspace size (GB)
nms=False, # TF: add NMS to model
agnostic_nms=False, # TF: add agnostic NMS to model
topk_per_class=100, # TF.js NMS: topk per class to keep
topk_all=100, # TF.js NMS: topk for all classes to keep
iou_thres=0.45, # TF.js NMS: IoU threshold
conf_thres=0.25, # TF.js NMS: confidence threshold
)
| 489 | |
| 490 | @smart_inference_mode() |
| 491 | def run( |
| 492 | data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' |
| 493 | weights=ROOT / 'yolov5s.pt', # weights path |
| 494 | imgsz=(640, 640), # image (height, width) |
| 495 | batch_size=1, # batch size |
| 496 | device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu |
| 497 | include=('torchscript', 'onnx'), # include formats |
| 498 | half=False, # FP16 half-precision export |
| 499 | inplace=False, # set YOLOv5 Detect() inplace=True |
| 500 | keras=False, # use Keras |
| 501 | optimize=False, # TorchScript: optimize for mobile |
| 502 | int8=False, # CoreML/TF INT8 quantization |
| 503 | dynamic=False, # ONNX/TF/TensorRT: dynamic axes |
| 504 | simplify=False, # ONNX: simplify model |
| 505 | opset=12, # ONNX: opset version |
| 506 | verbose=False, # TensorRT: verbose log |
| 507 | workspace=4, # TensorRT: workspace size (GB) |
| 508 | nms=False, # TF: add NMS to model |
| 509 | agnostic_nms=False, # TF: add agnostic NMS to model |
| 510 | topk_per_class=100, # TF.js NMS: topk per class to keep |
| 511 | topk_all=100, # TF.js NMS: topk for all classes to keep |
| 512 | iou_thres=0.45, # TF.js NMS: IoU threshold |
| 513 | conf_thres=0.25, # TF.js NMS: confidence threshold |
| 514 | ): |
| 515 | t = time.time() |
| 516 | include = [x.lower() for x in include] # to lowercase |
| 517 | fmts = tuple(export_formats()['Argument'][1:]) # --include arguments |
| 518 | flags = [x in include for x in fmts] |
| 519 | assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}' |
| 520 | jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans |
| 521 | file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) # PyTorch weights |
| 522 | |
| 523 | # Load PyTorch model |
| 524 | device = select_device(device) |
| 525 | if half: |
| 526 | assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0' |
| 527 | assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both' |
| 528 | model = attempt_load(weights, device=device, inplace=True, fuse=True) # load FP32 model |
| 529 | |
| 530 | # Checks |
| 531 | imgsz *= 2 if len(imgsz) == 1 else 1 # expand |
| 532 | if optimize: |
| 533 | assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' |
| 534 | |
| 535 | # Input |
| 536 | gs = int(max(model.stride)) # grid size (max stride) |
| 537 | imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples |
| 538 | im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection |
| 539 | |
| 540 | # Update model |
| 541 | model.eval() |
| 542 | for k, m in model.named_modules(): |
| 543 | if isinstance(m, Detect): |
| 544 | m.inplace = inplace |
| 545 | m.dynamic = dynamic |
| 546 | m.export = True |
| 547 | |
| 548 | for _ in range(2): |
no test coverage detected