集智書童 - YOLOV5 v6.1更新 | TensorRT+TPU+OpenVINO+TFJS+TFLite等平台一鍵導出和部署－鑽石舞台

Feb 23 Wed 2022 13:31
集智書童 - YOLOV5 v6.1更新 | TensorRT+TPU+OpenVINO+TFJS+TFLite等平台一鍵導出和部署

yolov5 release 6.1版本增加了TensorRT、Edge TPU和OpenVINO的支持，並提供了新的默認單周期線性LR調度器，以128批處理大小的再訓練模型。YOLOv5現在正式支持11種不同的權重，不僅可以直接導出，還可以用於推理(detect.py和PyTorch Hub)，以及在導出後對mAP配置文件和速度結果進行驗證。

比如，onnx文件的導出：

onnx導出1重大更新

TensorRT支持:TensorFlow, Keras, TFLite, TF.js模型導出現在完全集成使用python export.py -include saved_model pb TFLite tfjs

TensorFlow Edge TPU:新的更小的YOLOv5n(1.9M params)模型低於YOLOv5s(7.5M params)，導出到2.1 MB INT8大小，理想的超輕邊緣解決方案。

OpenVINO支持:YOLOv5 ONNX模型現在兼容OpenCV DNN和ONNX運行。

Export Benchmarks:使用python utils/ Benchmark.py導出所有YOLOv5格式(mAP和速度)。目前在CPU上運行，未來的更新將實現GPU支持。

架構:無更改。

超參數:小更改。Yaml LRF從0.2降至0.1。

訓練:默認學習速率(LR)調度器更新了一個周期的餘弦替換為一個周期的線性，以改善結果。

新版模型導出1、onnxdefexport_onnx(model,im,file,opset,train,dynamic,simplify,prefix=colorstr('ONNX:')):#YOLOv5ONNXexporttry:check_requirements(('onnx',))importonnxLOGGER.info(f'\n{prefix}startingexportwithonnx{onnx.__version__}...')f=file.with_suffix('.onnx')torch.onnx.export(model,im,f,verbose=False,opset_version=opset,training=torch.onnx.TrainingMode.TRAININGiftrainelsetorch.onnx.TrainingMode.EVAL,do_constant_folding=nottrain,input_names=['images'],output_names=['output'],dynamic_axes={'images':{0:'batch',2:'height',3:'width'},#shape(1,3,640,640)'output':{0:'batch',1:'anchors'}#shape(1,25200,85)}ifdynamicelseNone)#Checksmodel_onnx=onnx.load(f)#loadonnxmodelonnx.checker.check_model(model_onnx)#checkonnxmodel#LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))#print#Simplifyifsimplify:try:check_requirements(('onnx-simplifier',))importonnxsimLOGGER.info(f'{prefix}simplifyingwithonnx-simplifier{onnxsim.__version__}...')model_onnx,check=onnxsim.simplify(model_onnx,dynamic_input_shape=dynamic,input_shapes={'images':list(im.shape)}ifdynamicelseNone)assertcheck,'assertcheckfailed'onnx.save(model_onnx,f)exceptExceptionase:LOGGER.info(f'{prefix}simplifierfailure:{e}')LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'{prefix}exportfailure:{e}')2、 openvinodefexport_openvino(model,im,file,prefix=colorstr('OpenVINO:')):#YOLOv5OpenVINOexporttry:check_requirements(('openvino-dev',))#requiresopenvino-dev:https://pypi.org/project/openvino-dev/importopenvino.inference_engineasieLOGGER.info(f'\n{prefix}startingexportwithopenvino{ie.__version__}...')f=str(file).replace('.pt','_openvino_model'+os.sep)cmd=f"mo--input_model{file.with_suffix('.onnx')}--output_dir{f}"subprocess.check_output(cmd,shell=True)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')3、 coremldefexport_coreml(model,im,file,prefix=colorstr('CoreML:')):#YOLOv5CoreMLexporttry:check_requirements(('coremltools',))importcoremltoolsasctLOGGER.info(f'\n{prefix}startingexportwithcoremltools{ct.__version__}...')f=file.with_suffix('.mlmodel')ts=torch.jit.trace(model,im,strict=False)#TorchScriptmodelct_model=ct.convert(ts,inputs=[ct.ImageType('image',shape=im.shape,scale=1/255,bias=[0,0,0])])ct_model.save(f)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnct_model,fexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')returnNone,None4、TensorRTdefexport_engine(model,im,file,train,half,simplify,workspace=4,verbose=False,prefix=colorstr('TensorRT:')):#YOLOv5TensorRTexporthttps://developer.nvidia.com/tensorrttry:check_requirements(('tensorrt',))importtensorrtastrtiftrt.__version__[0]=='7':#TensorRT7handlinghttps://github.com/ultralytics/yolov5/issues/6012grid=model.model[-1].anchor_gridmodel.model[-1].anchor_grid=[a[...,:1,:1,:]foraingrid]export_onnx(model,im,file,12,train,False,simplify)#opset12model.model[-1].anchor_grid=gridelse:#TensorRT>=8check_version(trt.__version__,'8.0.0',hard=True)#requiretensorrt>=8.0.0export_onnx(model,im,file,13,train,False,simplify)#opset13onnx=file.with_suffix('.onnx')LOGGER.info(f'\n{prefix}startingexportwithTensorRT{trt.__version__}...')assertim.device.type!='cpu','exportrunningonCPUbutmustbeonGPU,i.e.`pythonexport.py--device0`'assertonnx.exists(),f'failedtoexportONNXfile:{onnx}'f=file.with_suffix('.engine')#TensorRTenginefilelogger=trt.Logger(trt.Logger.INFO)ifverbose:logger.min_severity=trt.Logger.Severity.VERBOSEbuilder=trt.Builder(logger)config=builder.create_builder_config()config.max_workspace_size=workspace*1<<30flag=(1<<int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))network=builder.create_network(flag)parser=trt.OnnxParser(network,logger)ifnotparser.parse_from_file(str(onnx)):raiseRuntimeError(f'failedtoloadONNXfile:{onnx}')inputs=[network.get_input(i)foriinrange(network.num_inputs)]outputs=[network.get_output(i)foriinrange(network.num_outputs)]LOGGER.info(f'{prefix}NetworkDescription:')forinpininputs:LOGGER.info(f'{prefix}\tinput"{inp.name}"withshape{inp.shape}anddtype{inp.dtype}')foroutinoutputs:LOGGER.info(f'{prefix}\toutput"{out.name}"withshape{out.shape}anddtype{out.dtype}')half&=builder.platform_has_fast_fp16LOGGER.info(f'{prefix}buildingFP{16ifhalfelse32}enginein{f}')ifhalf:config.set_flag(trt.BuilderFlag.FP16)withbuilder.build_engine(network,config)asengine,open(f,'wb')ast:t.write(engine.serialize())LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')5、TensorFlowdefexport_pb(keras_model,im,file,prefix=colorstr('TensorFlowGraphDef:')):#YOLOv5TensorFlowGraphDef*.pbexporthttps://github.com/leimao/Frozen_Graph_TensorFlowtry:importtensorflowastffromtensorflow.python.framework.convert_to_constantsimportconvert_variables_to_constants_v2LOGGER.info(f'\n{prefix}startingexportwithtensorflow{tf.__version__}...')f=file.with_suffix('.pb')m=tf.function(lambdax:keras_model(x))#fullmodelm=m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape,keras_model.inputs[0].dtype))frozen_func=convert_variables_to_constants_v2(m)frozen_func.graph.as_graph_def()tf.io.write_graph(graph_or_graph_def=frozen_func.graph,logdir=str(f.parent),name=f.name,as_text=False)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')6、TensorFlow-Litedefexport_tflite(keras_model,im,file,int8,data,ncalib,prefix=colorstr('TensorFlowLite:')):#YOLOv5TensorFlowLiteexporttry:importtensorflowastfLOGGER.info(f'\n{prefix}startingexportwithtensorflow{tf.__version__}...')batch_size,ch,*imgsz=list(im.shape)#BCHWf=str(file).replace('.pt','-fp16.tflite')converter=tf.lite.TFLiteConverter.from_keras_model(keras_model)converter.target_spec.supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS]converter.target_spec.supported_types=[tf.float16]converter.optimizations=[tf.lite.Optimize.DEFAULT]ifint8:frommodels.tfimportrepresentative_dataset_gendataset=LoadImages(check_dataset(data)['train'],img_size=imgsz,auto=False)#representativedataconverter.representative_dataset=lambda:representative_dataset_gen(dataset,ncalib)converter.target_spec.supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS_INT8]converter.target_spec.supported_types=[]converter.inference_input_type=tf.uint8#ortf.int8converter.inference_output_type=tf.uint8#ortf.int8converter.experimental_new_quantizer=Falsef=str(file).replace('.pt','-int8.tflite')tflite_model=converter.convert()open(f,"wb").write(tflite_model)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')7、Egde TPUdefexport_edgetpu(keras_model,im,file,prefix=colorstr('EdgeTPU:')):#YOLOv5EdgeTPUexporthttps://coral.ai/docs/edgetpu/models-intro/try:cmd='edgetpu_compiler--version'help_url='https://coral.ai/docs/edgetpu/compiler/'assertplatform.system()=='Linux',f'exportonlysupportedonLinux.See{help_url}'ifsubprocess.run(cmd+'>/dev/null',shell=True).returncode!=0:LOGGER.info(f'\n{prefix}exportrequiresEdgeTPUcompiler.Attemptinginstallfrom{help_url}')sudo=subprocess.run('sudo--version>/dev/null',shell=True).returncode==0#sudoinstalledonsystemforcin['curlhttps://packages.cloud.google.com/apt/doc/apt-key.gpg|sudoapt-keyadd-','echo"debhttps://packages.cloud.google.com/aptcoral-edgetpu-stablemain"|sudotee/etc/apt/sources.list.d/coral-edgetpu.list','sudoapt-getupdate','sudoapt-getinstalledgetpu-compiler']:subprocess.run(cifsudoelsec.replace('sudo',''),shell=True,check=True)ver=subprocess.run(cmd,shell=True,capture_output=True,check=True).stdout.decode().split()[-1]LOGGER.info(f'\n{prefix}startingexportwithEdgeTPUcompiler{ver}...')f=str(file).replace('.pt','-int8_edgetpu.tflite')#EdgeTPUmodelf_tfl=str(file).replace('.pt','-int8.tflite')#TFLitemodelcmd=f"edgetpu_compiler-s{f_tfl}"subprocess.run(cmd,shell=True,check=True)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')8、TensorFlow.jsdefexport_tfjs(keras_model,im,file,prefix=colorstr('TensorFlow.js:')):#YOLOv5TensorFlow.jsexporttry:check_requirements(('tensorflowjs',))importreimporttensorflowjsastfjsLOGGER.info(f'\n{prefix}startingexportwithtensorflowjs{tfjs.__version__}...')f=str(file).replace('.pt','_web_model')#jsdirf_pb=file.with_suffix('.pb')#*.pbpathf_json=f+'/model.json'#*.jsonpathcmd=f'tensorflowjs_converter--input_format=tf_frozen_model'\f'--output_node_names="Identity,Identity_1,Identity_2,Identity_3"{f_pb}{f}'subprocess.run(cmd,shell=True)json=open(f_json).read()withopen(f_json,'w')asj:#sortJSONIdentity_*inascendingordersubst=re.sub(r'{"outputs":{"Identity.?.?":{"name":"Identity.?.?"},'r'"Identity.?.?":{"name":"Identity.?.?"},'r'"Identity.?.?":{"name":"Identity.?.?"},'r'"Identity.?.?":{"name":"Identity.?.?"}}}',r'{"outputs":{"Identity":{"name":"Identity"},'r'"Identity_1":{"name":"Identity_1"},'r'"Identity_2":{"name":"Identity_2"},'r'"Identity_3":{"name":"Identity_3"}}}',json)j.write(subst)LOGGER.info(f'{prefix}exportsuccess,savedas{f}({file_size(f):.1f}MB)')returnfexceptExceptionase:LOGGER.info(f'\n{prefix}exportfailure:{e}')新版檢測推理pythonpath/to/detect.py--weightsyolov5s.pt#PyTorchyolov5s.torchscript#TorchScriptyolov5s.onnx#ONNXRuntimeorOpenCVDNNwith--dnnyolov5s.xml#OpenVINOyolov5s.engine#TensorRTyolov5s.mlmodel#CoreML(MacOS-only)yolov5s_saved_model#TensorFlowSavedModelyolov5s.pb#TensorFlowGraphDefyolov5s.tflite#TensorFlowLiteyolov5s_edgetpu.tflite#TensorFlowEdgeTPU2最新結果