yolov5 Coremltools

在部署yolov5时如果在ios端选择CoreML作为部署框架，直接导出包含后处理部分的.mlmodel会非常方便，但是在YOLOv5官方代码中提供的export.py文件导出的文件是不包含后处理部分的，感谢大神的奉献，我在源代码的基础上做了少量调整，便于导出时设置各种参数：

import torch
import coremltools as ct
from pathlib import Path# Add silu function for yolov5s v4 model: https://github.com/apple/coremltools/issues/1099
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil import register_torch_op
from coremltools.converters.mil.frontend.torch.ops import _get_inputs# 自定义激活函数，因为在coremltools中不支持此操作
@register_torch_op
def silu(context, node):inputs = _get_inputs(context, node, expected=1)x = inputs[0]y = mb.sigmoid(x=x)z = mb.mul(x=x, y=y, name=node.name)context.add(z)# 特征图大小，默认有三组，因为有三个下采样的步长
def featureMapDimensions(imgSize, strides):return [imgSize // stride for stride in strides]def make_grid(nx, ny):yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])return torch.stack((xv, yv), 2).view((ny, nx, 2)).float()# 导出torchscript文件
def exportTorchscript(model, sampleInput, checkInputs, fileName):'''Traces a pytorch model and produces a TorchScript'''try:print(f'Starting TorchScript export with torch {torch.__version__}')ts = torch.jit.trace(model, sampleInput, check_inputs=checkInputs)ts.save(fileName)print(f'TorchScript export success, saved as {fileName}')return tsexcept Exception as e:print(f'TorchScript export failure: {e}')# 获得coremltools可以操作的spec
def convertToCoremlSpec(torchScript, sampleInput):'''Converts a torchscript to a coreml model'''try:print(f'Starting CoreML conversion with coremltoola {ct.__version__}')nnSpec = ct.convert(torchScript, inputs=[ct.ImageType(name='image', shape=sampleInput.shape, scale=1 / 255.0, bias=[0, 0, 0])]).get_spec()print(f'CoreML conversion success')except Exception as e:print(f'CoreML conversion failure: {e}')returnreturn nnSpec# 设置输出的shape和数据类型
def addOutputMetaData(nnSpec,outputSize,featureMapDimensionss):'''Adds the correct output shapes and data types to the coreml model'''for i, featureMapDimension in enumerate(featureMapDimensionss):nnSpec.description.output[i].type.multiArrayType.shape.append(1)nnSpec.description.output[i].type.multiArrayType.shape.append(3)nnSpec.description.output[i].type.multiArrayType.shape.append(featureMapDimension)nnSpec.description.output[i].type.multiArrayType.shape.append(featureMapDimension)# pc, bx, by, bh, bw, c (no of class class labels)nnSpec.description.output[i].type.multiArrayType.shape.append(outputSize)nnSpec.description.output[i].type.multiArrayType.dataType = ct.proto.FeatureTypes_pb2.ArrayFeatureType.DOUBLE# 添加输出层
def addExportLayerToCoreml(builder, numberOfClassLabels, imgSize, featureMapDimensionss, anchorGrid):'''Adds the yolov5 export layer to the coreml model'''outputNames = [output.name for output in builder.spec.description.output]for i, outputName in enumerate(outputNames):# formulas: https://github.com/ultralytics/yolov5/issues/471builder.add_activation(name=f"sigmoid_{outputName}", non_linearity="SIGMOID",input_name=outputName, output_name=f"{outputName}_sigmoid")### Coordinates calculation #### input (1, 3, nC, nC, cls+5), output (1, 3, nC, nC, 2) -> nC = 640 / strides[i]   cls为总的类别数量builder.add_slice(name=f"slice_coordinates_xy_{outputName}", input_name=f"{outputName}_sigmoid",output_name=f"{outputName}_sliced_coordinates_xy", axis="width", start_index=0, end_index=2)# x,y * 2builder.add_elementwise(name=f"multiply_xy_by_two_{outputName}", input_names=[f"{outputName}_sliced_coordinates_xy"], output_name=f"{outputName}_multiplied_xy_by_two", mode="MULTIPLY", alpha=2)# x,y * 2 - 0.5builder.add_elementwise(name=f"subtract_0_5_from_xy_{outputName}", input_names=[f"{outputName}_multiplied_xy_by_two"], output_name=f"{outputName}_subtracted_0_5_from_xy", mode="ADD", alpha=-0.5)grid = make_grid(featureMapDimensionss[i], featureMapDimensionss[i]).numpy()# x,y * 2 - 0.5 + grid[i]builder.add_bias(name=f"add_grid_from_xy_{outputName}", input_name=f"{outputName}_subtracted_0_5_from_xy",output_name=f"{outputName}_added_grid_xy", b=grid, shape_bias=grid.shape)# (x,y * 2 - 0.5 + grid[i]) * stride[i]builder.add_elementwise(name=f"multiply_xy_by_stride_{outputName}", input_names=[f"{outputName}_added_grid_xy"], output_name=f"{outputName}_calculated_xy", mode="MULTIPLY", alpha=strides[i])# input (1, 3, nC, nC, cls+5), output (1, 3, nC, nC, 2)builder.add_slice(name=f"slice_coordinates_wh_{outputName}", input_name=f"{outputName}_sigmoid",output_name=f"{outputName}_sliced_coordinates_wh", axis="width", start_index=2, end_index=4)# w,h * 2builder.add_elementwise(name=f"multiply_wh_by_two_{outputName}", input_names=[f"{outputName}_sliced_coordinates_wh"], output_name=f"{outputName}_multiplied_wh_by_two", mode="MULTIPLY", alpha=2)# (w,h * 2) ** 2builder.add_unary(name=f"power_wh_{outputName}", input_name=f"{outputName}_multiplied_wh_by_two",output_name=f"{outputName}_power_wh", mode="power", alpha=2)# (w,h * 2) ** 2 * anchor_grid[i]anchor = anchorGrid[i].expand(-1, featureMapDimensionss[i],featureMapDimensionss[i], -1).numpy()builder.add_load_constant_nd(name=f"anchors_{outputName}", output_name=f"{outputName}_anchors", constant_value=anchor, shape=anchor.shape)builder.add_elementwise(name=f"multiply_wh_with_achors_{outputName}", input_names=[f"{outputName}_power_wh", f"{outputName}_anchors"], output_name=f"{outputName}_calculated_wh", mode="MULTIPLY")builder.add_concat_nd(name=f"concat_coordinates_{outputName}", input_names=[f"{outputName}_calculated_xy", f"{outputName}_calculated_wh"], output_name=f"{outputName}_raw_coordinates", axis=-1)builder.add_scale(name=f"normalize_coordinates_{outputName}", input_name=f"{outputName}_raw_coordinates",output_name=f"{outputName}_raw_normalized_coordinates", W=torch.tensor([1 / imgSize]).numpy(), b=0, has_bias=False)### Confidence calculation ###builder.add_slice(name=f"slice_object_confidence_{outputName}", input_name=f"{outputName}_sigmoid",output_name=f"{outputName}_object_confidence", axis="width", start_index=4, end_index=5)builder.add_slice(name=f"slice_label_confidence_{outputName}", input_name=f"{outputName}_sigmoid",output_name=f"{outputName}_label_confidence", axis="width", start_index=5, end_index=0)# confidence = object_confidence * label_confidencebuilder.add_multiply_broadcastable(name=f"multiply_object_label_confidence_{outputName}", input_names=[f"{outputName}_label_confidence", f"{outputName}_object_confidence"], output_name=f"{outputName}_raw_confidence")# input: (1, 3, nC, nC, cls+5), output: (3 * nc^2, cls+5)builder.add_flatten_to_2d(name=f"flatten_confidence_{outputName}", input_name=f"{outputName}_raw_confidence", output_name=f"{outputName}_flatten_raw_confidence", axis=-1)builder.add_flatten_to_2d(name=f"flatten_coordinates_{outputName}", input_name=f"{outputName}_raw_normalized_coordinates", output_name=f"{outputName}_flatten_raw_coordinates", axis=-1)builder.add_concat_nd(name="concat_confidence", input_names=[f"{outputName}_flatten_raw_confidence" for outputName in outputNames], output_name="raw_confidence", axis=-2)builder.add_concat_nd(name="concat_coordinates", input_names=[f"{outputName}_flatten_raw_coordinates" for outputName in outputNames], output_name="raw_coordinates", axis=-2)builder.set_output(output_names=["raw_confidence", "raw_coordinates"], output_dims=[(int(3*((imgSize/strides[0])**2+(imgSize/strides[1])**2+(imgSize/strides[2])**2)),numberOfClassLabels),(int(3*((imgSize/strides[0])**2+(imgSize/strides[1])**2+(imgSize/strides[2])**2)), 4)])#　添加设置nms
def createNmsModelSpec(nnSpec, numberOfClassLabels, classLabels):'''Create a coreml model with nms to filter the results of the model'''nmsSpec = ct.proto.Model_pb2.Model()nmsSpec.specificationVersion = 4# Define input and outputs of the modelfor i in range(2):nnOutput = nnSpec.description.output[i].SerializeToString()nmsSpec.description.input.add()nmsSpec.description.input[i].ParseFromString(nnOutput)nmsSpec.description.output.add()nmsSpec.description.output[i].ParseFromString(nnOutput)nmsSpec.description.output[0].name = "confidence"nmsSpec.description.output[1].name = "coordinates"# Define output shape of the modeloutputSizes = [numberOfClassLabels, 4]for i in range(len(outputSizes)):maType = nmsSpec.description.output[i].type.multiArrayType# First dimension of both output is the number of boxes, which should be flexiblemaType.shapeRange.sizeRanges.add()maType.shapeRange.sizeRanges[0].lowerBound = 0maType.shapeRange.sizeRanges[0].upperBound = -1# Second dimension is fixed, for "confidence" it's the number of classes, for coordinates it's position (x, y) and size (w, h)maType.shapeRange.sizeRanges.add()maType.shapeRange.sizeRanges[1].lowerBound = outputSizes[i]maType.shapeRange.sizeRanges[1].upperBound = outputSizes[i]del maType.shape[:]# Define the model type non maximum supressionnms = nmsSpec.nonMaximumSuppressionnms.confidenceInputFeatureName = "raw_confidence"nms.coordinatesInputFeatureName = "raw_coordinates"nms.confidenceOutputFeatureName = "confidence"nms.coordinatesOutputFeatureName = "coordinates"nms.iouThresholdInputFeatureName = "iouThreshold"nms.confidenceThresholdInputFeatureName = "confidenceThreshold"# Some good default values for the two additional inputs, can be overwritten when using the modelnms.iouThreshold = 0.6nms.confidenceThreshold = 0.4nms.stringClassLabels.vector.extend(classLabels)return nmsSpec# 导出mlmodel
def combineModelsAndExport(builderSpec, nmsSpec, fileName, imgSize,quantize):'''Combines the coreml model with export logic and the nms to one final model. Optionally save with different quantization (32, 16, 8) (Works only if on Mac Os)'''try:print(f'Combine CoreMl model with nms and export model')# Combine models to a single onepipeline = ct.models.pipeline.Pipeline(input_features=[("image", ct.models.datatypes.Array(3, imgSize, imgSize)),("iouThreshold", ct.models.datatypes.Double()),("confidenceThreshold", ct.models.datatypes.Double())], output_features=["confidence", "coordinates"])# Required version (>= ios13) in order for mns to workpipeline.spec.specificationVersion = 4pipeline.add_model(builderSpec)pipeline.add_model(nmsSpec)pipeline.spec.description.input[0].ParseFromString(builderSpec.description.input[0].SerializeToString())pipeline.spec.description.output[0].ParseFromString(nmsSpec.description.output[0].SerializeToString())pipeline.spec.description.output[1].ParseFromString(nmsSpec.description.output[1].SerializeToString())# Metadata for the model‚pipeline.spec.description.input[1].shortDescription = "(optional) IOU Threshold override (Default: 0.6)"pipeline.spec.description.input[2].shortDescription = "(optional) Confidence Threshold override (Default: 0.4)"pipeline.spec.description.output[0].shortDescription = u"Boxes \xd7 Class confidence"pipeline.spec.description.output[1].shortDescription = u"Boxes \xd7 [x, y, width, height] (relative to image size)"pipeline.spec.description.metadata.versionString = "1.0"pipeline.spec.description.metadata.shortDescription = "yolov5"pipeline.spec.description.metadata.author = "Leon De Andrade"pipeline.spec.description.metadata.license = ""model = ct.models.MLModel(pipeline.spec)model.save(fileName)if quantize:fileName16 = fileName.replace(".mlmodel", "_16.mlmodel")modelFp16 = ct.models.neural_network.quantization_utils.quantize_weights(model, nbits=16)modelFp16.save(fileName16)fileName8 = fileName.replace(".mlmodel", "_8.mlmodel")modelFp8 = ct.models.neural_network.quantization_utils.quantize_weights(model, nbits=8)modelFp8.save(fileName8)print(f'CoreML export success, saved as {fileName}')except Exception as e:print(f'CoreML export failure: {e}')# 是否对步长和anchor顺序反转
def reversemodel(reverse, strides, anchors):if reverse:strides.reverse()anchors = anchors[::-1]return strides, anchorsdef main(strides, anchors, reverseModel=False):strides,anchors = reversemodel(reverseModel, strides, anchors)anchorGrid = torch.tensor(anchors).float().view(3, -1, 1, 1, 2)numberOfClassLabels = len(classLabels)outputSize = numberOfClassLabels + 5featureMapDimensionss = featureMapDimensions(imgSize, strides)if not Path(model_input_path).exists():print("Error: Input model not found")returnPath(model_output_directory).mkdir(parents=True, exist_ok=True)# 输入的shapesampleInput = torch.zeros((1, 3, imgSize, imgSize))checkInputs = [(torch.rand(1, 3, imgSize, imgSize),),(torch.rand(1, 3, imgSize, imgSize),)]# 加载pt文件model = torch.load(model_input_path, map_location=torch.device('cpu'))['model'].float()model.eval()model.model[-1].export = True# Dry run, necessary for correct tracing!model(sampleInput)# 导出为torchscriptts = exportTorchscript(model, sampleInput, checkInputs,f"{model_output_directory}/{model_output_name}.torchscript.pt")# 获得coremltools可以操纵的spec# Convert pytorch to raw coreml modelmodelSpec = convertToCoremlSpec(ts, sampleInput)addOutputMetaData(modelSpec, outputSize,featureMapDimensionss)# Add export logic to coreml modelbuilder = ct.models.neural_network.NeuralNetworkBuilder(spec=modelSpec)addExportLayerToCoreml(builder, numberOfClassLabels,imgSize,featureMapDimensionss, anchorGrid)# Create nms logicnmsSpec = createNmsModelSpec(builder.spec, numberOfClassLabels, classLabels)# Combine model with export logic and nms logiccombineModelsAndExport(builder.spec, nmsSpec, f"{model_output_directory}/{model_output_name}.mlmodel", imgSize, quantize)if __name__ == '__main__':# ------------------------------------------------------------------------------## ---------------------------------需要自己修改的常用参数--------------------------## 图片的大小imgSize = 641# 标签的名称classLabels = ['a', 'b', 'c', 'd', 'e', '1', '2', '3', '4']  # 类别名称# 需要转换的pt文件model_input_path = './best.pt'# mlmodel输出的文件夹路径model_output_directory = "./"# 输出的mlmodel的命名(.mlmodel的前缀名称)model_output_name = 'yolov5-best-iOS'# 下采样的步长strides = [8, 16, 32]# anchor,与yolov5yaml中的文件相同anchors = ([10,13, 16,30, 33,23], [30,61, 62,45, 59,119], [116,90, 156,198, 373,326]) # 步长的顺序是否相反modelreverse = True# 是否对mlmodel进行量化quantize = False   # WARNING! Unable to return a quantized MLModel instance sinceOS != macOS 10.14 or later# ---------------------------------需要自己修改的常用参数（完）--------------------------## ----------------------------------------------------------------------------------- #main(strides,anchors,modelreverse)

比较：

使用yolov5官方的export.py导出的.mlmodel文件，使用本文件导出的.mlmodel：

yolov5 Coremltools相关推荐

用 YOLOv5模型识别出表情！
作者 | 闫永强来源 | Datawhale 本文利用YOLOV5对手势进行训练识别,并识别显示出对应的emoji,如同下图: 本文整体思路如下.提示:本文含完整实践代码,代码较长,建议先看文字部分 ...
用YOLOv5模型识别出表情！
↑↑↑关注后"星标"Datawhale 每日干货 & 每月组队学习,不错过 Datawhale干货作者:闫永强,算法工程师,Datawhale成员本文利用YOLOV5对 ...
干货 | YOLOV5 训练自动驾驶数据集，并转Tensorrt，收藏！
点击上方"小白学视觉",选择加"星标"或"置顶" 重磅干货,第一时间送达本文转自|AI算法与图像处理准备数据集环境配置配置文件修改 ...
YOLOV5 v6.1更新 | TensorRT+TPU+OpenVINO+TFJS+TFLite等平台一键导出和部署
点击上方"3D视觉工坊",选择"星标" 干货第一时间送达作者丨ChaucerG 来源丨集智书童 yolov5 release 6.1版本增加了TensorRT ...
YOLOv5导出jit，onnx，engine
一.YOLOv5导出jit YOLOv5自导出,我们可以直接用它的导出代码:models/export.py """Exports a YOLOv5 *.pt model ...
YOLOv5的pytorch模型文件转换为ONNX文件
YOLOv5 YOLOv5下载与测试运行导出ONNX格式文件 ONNX转为为IR中间格式环境: Windows 10 Anaconda 2.0.4 OpenVINO 工具包 2021.2 Pyth ...
(五)将YOLOv5 PyTorch模型权重转换为TensorFlow Lite格式
目录介绍建议使用Ultralytics进行重量转换的方法在Google Colab上测试TFLite权重 TFLite解释器,可在边缘设备上实现良好性能在本地测试YOLOv5模型权重下一步 ...
适用于openvino 2020.2的yolov5的docker制作
docker版本:20.0 系统环境:ubuntu16.04 目标openvino版本:2020.2 这里主要提供yolov5的docker制作过程,在该docker中,训练得到的.pt模型将适用于o ...
win7下用docker部署的基于openvino的yolov5算法（三）yolov5 v4.0环境安装以及.pt转成.ir模型
在上一篇博文中,我们创建了基础的openvino容器openvino2020.3.1_ubuntu18.04,具体环境为 * ubuntu 18.04 * openvino2020.3.341 这里, ...
Windows下部署yolov5实现口罩检测
目录前言一.部署yoloV5 1.安装依赖二.使用步骤 1.准备自己的数据集 2.更改配置 3.开始训练 4.可视化总结前言本人环境声明: 系统环境:WIN10 cuda版本:10.2.8 ...

yolov5 Coremltools

yolov5 Coremltools相关推荐

最新文章

热门文章