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、onnx

defexport_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、 openvino

defexport_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、 coreml

defexport_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、TensorRT

defexport_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<5、TensorFlow

defexport_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-Lite

defexport_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 TPU

defexport_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.js

defexport_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最新结果3与6.0版本的精度对比4参考

[1].https://github.com/ultralytics/yolov5

5推荐阅读

往期推荐

超越ConvNeXt | 大道至简,VAN用普通卷积,登顶Backbone性能巅峰(附代码解读)

超级干货 | 用万字文章总结25种正则化方法(值得收藏)

小目标Trick | Detectron2、MMDetection、YOLOv5都通用的小目标检测解决方案

ShiftViT用Swin Transformer的精度跑赢ResNet的速度,论述ViT的成功不在注意力!

阿里提出QuadTree Transformer | 最轻、最强的Vision Transformer Backbone

MoA-Transformer | Swin-Transformer应该如何更好地引入全局信息?

致敬ATSS | Dynamic ATSS再造ATSS辉煌!!!

长按扫描下方二维码添加小助手并加入交流群,群里博士大佬云集,每日讨论话题有目标检测、语义分割、超分辨率、模型部署、数学基础知识、算法面试题分享的等等内容,当然也少不了搬砖人的扯犊子

长按扫描下方二维码添加小助手。

可以一起讨论遇到的问题

声明:转载请说明出处

扫描下方二维码关注【集智书童】公众号,获取更多实践项目源码和论文解读,非常期待你我的相遇,让我们以梦为马,砥砺前行!

推荐内容