65.9K
CodeProject 正在变化。 阅读更多。
Home

创建带解码逻辑的 YOLO Core ML 对象检测器

Jarek Szczegielniak
starIconstarIconstarIconstarIconstarIcon

5.00/5 (2投票s)

2020年11月25日

CPOL

4分钟阅读

viewsIcon

7148

downloadIcon

108

在本文中，我们准备直接在 Core ML 模型中包含检测解码。

引言

本系列假设您熟悉 Python、Conda 和 ONNX，并且具有在 Xcode 中开发 iOS 应用程序的经验。欢迎您下载此项目的源代码。我们将使用 macOS 10.15+、Xcode 11.7+ 和 iOS 13+ 运行代码。

缩小模型

为了在 iOS 设备上节省内存，而不会对模型的性能产生负面影响，我们应该将其权重从 32 位精度降低到 16 位精度。请注意，当模型在 iOS 设备的 GPU 或神经引擎上执行时（它应该这样做），它始终以 16 位浮点数运行。只有在 CPU 上运行时，32 位精度才会有所不同。

让我们开始吧

import os
import coremltools as ct
import numpy as np

model_converted = ct.models.MLModel('./models/yolov2-coco-9.mlmodel')
model_converted = ct.models.neural_network.quantization_utils.quantize_weights(
    model_converted, 
    nbits=16, 
    quantization_mode='linear')
model_converted.save('./models/yolov2-16.mlmodel')

构建 YOLO 解码器

我们有两个选择：将解码器层添加到现有模型中，或者创建一个单独的模型，然后使用管道连接这两个模型。让我们选择后一种方案。

我们将从创建一个新的 NeuralNetworkBuilder 实例并映射新解码器模型的输入和输出开始

from coremltools.models import datatypes

input_features = [ (spec.description.output[0].name, datatypes.Array(1, 425, 13, 13)) ]
output_features = [ ('all_scores', datatypes.Array(1, 845, 80)),
                    ('all_boxes', datatypes.Array(1, 845, 4)) ]

builder = ct.models.neural_network.NeuralNetworkBuilder(
    input_features, 
    output_features, 
    disable_rank5_shape_mapping=True
)

builder.spec.description.input[0].ParseFromString(spec.description.output[0].SerializeToString())

接下来，我们定义计算所需的常量

GRID_SIZE = 13
CELL_SIZE = 1 / GRID_SIZE 
BOXES_PER_CELL = 5
NUM_CLASSES = 80

ANCHORS_W = np.array([0.57273, 1.87446, 3.33843, 7.88282, 9.77052]).reshape(1, 1, 5)
ANCHORS_H = np.array([0.677385, 2.06253, 5.47434, 3.52778, 9.16828]).reshape(1, 1, 5)

CX = np.tile(np.arange(GRID_SIZE), GRID_SIZE).reshape(1, 1, GRID_SIZE**2, 1)
CY = np.tile(np.arange(GRID_SIZE), GRID_SIZE).reshape(1, GRID_SIZE, GRID_SIZE).transpose()
CY = CY.reshape(1, 1, GRID_SIZE**2, 1)

请注意上面的 CELL_SIZE 值。要将我们的模型与 Vision 框架一起使用，我们需要将边界框坐标从图像像素缩放到 [0-1] 范围。

要使用定义的常量进行计算，我们将它们添加到网络中

builder.add_load_constant_nd('CX', output_name='CX', constant_value=CX, shape=CX.shape)
builder.add_load_constant_nd('CY', output_name='CY', constant_value=CY, shape=CY.shape)
builder.add_load_constant_nd('ANCHORS_W', output_name='ANCHORS_W', constant_value=ANCHORS_W, shape=ANCHORS_W.shape)
builder.add_load_constant_nd('ANCHORS_H', output_name='ANCHORS_H', constant_value=ANCHORS_H, shape=ANCHORS_H.shape)

现在我们准备好将层添加到我们的 Core ML 模型中。在大多数情况下，这将是上一篇文章中代码的直接转换，并尽可能使用相同的变量/节点名称。但是，有时 Core ML 的一些怪癖会强制进行小的更改。请参阅代码下载以获取完整的解决方案，因为为了提高可读性，这里将不包括一些显而易见的代码序列。

我们从对应于上一篇文章中前两个转换的层开始（向量化实现）

builder.add_transpose(
    'yolo_trans_node', 
    axes=(0,2,3,1), 
    input_name='218', 
    output_name=‘yolo_transp')

builder.add_reshape_static(
    'yolo_reshap', 
    input_name='yolo_transp',
    output_name='yolo_reshap',
    output_shape=(1, GRID_SIZE**2, BOXES_PER_CELL, NUM_CLASSES + 5)
)

当我们使用 NeuralNetworkBuilder 实例创建一个新层时，我们需要为节点及其 output_name（在本例中分别为“yolo_trans_node”和“yolo_transp”）指定一个唯一的名称。input_name 值必须对应于现有的 output_name（在本例中为“218”，这是我们转换后的 YOLO v2 模型的输出）。

要提取编码的框和置信度值，我们需要拆分输入数组

builder.add_split_nd(
    'split_boxes_node', 
    input_name='yolo_reshap',
    output_names=['tx', 'ty', 'tw', 'th', 'tc', 'classes_raw'],    
    axis=3,
    split_sizes=[1, 1, 1, 1, 1, 80])

此操作将 raw_preds 数组分割成 tx、ty、tw、th、tc 和 classes_raw 数组，这些数组来自上一篇文章。

不幸的是，其余的代码将更加冗长，因为我们需要为每个基本算术运算使用一个单独的节点。这导致了以下情况，即我们向量化解码器中的简单一行代码

x = ((CX + sigmoid(tx)) * CELL_SIZE).reshape(-1)

变成

builder.add_reshape_static('tx:1', input_name='tx', output_name='tx:1', output_shape=(1,169,5))
builder.add_activation('tx:1_sigm', non_linearity='SIGMOID', input_name='tx:1', output_name='tx:1_sigm')
builder.add_add_broadcastable('tx:1_add', input_names=['CX', 'tx:1_sigm'], output_name='tx:1_add')
builder.add_elementwise('x', input_names=['tx:1_add'], output_name='x', mode='MULTIPLY', alpha=CELL_SIZE)

请注意，为了使代码更短更易读，我们在输出形状参数中使用显式值“169”而不是 GRID_SIZE**2，使用“5”而不是 BOXES_PER_CELL。这同样适用于在其他地方使用“80”而不是 NUM_CLASSES 字面量。当然，在一个正确且灵活的解决方案中，我们应该坚持使用字面量。

需要进行相同的运算来计算 y。然后我们有一个非常类似的代码来计算边界框宽度 (w)

builder.add_reshape_static('tw:1', input_name='tw', output_name='tw:1', output_shape=(1,169,5))
builder.add_unary('tw:1_exp', input_name='tw:1', output_name='tw:1_exp', mode='exp')
builder.add_multiply_broadcastable('tw:1_mul', input_names=['tw:1_exp', 'ANCHORS_W'], output_name='tw:1_mul')
builder.add_elementwise('w', input_names=['tw:1_mul'], output_name='w', mode='MULTIPLY', alpha=CELL_SIZE)

随后计算 h 也非常相似（除了使用 ANCHORS_H 而不是 ANCHORS_W 常量）。

最后，我们对 box_confidence 和 classes_confidence 值进行解码

builder.add_reshape_static('tc:1', input_name='tc', output_name='tc:1', output_shape=(1,169*5,1))
builder.add_activation('box_confidence', non_linearity='SIGMOID', input_name='tc:1', output_name='box_confidence')
builder.add_reshape_static('classes_raw:1', input_name='classes_raw', output_name='classes_raw:1', output_shape=(1,169*5,80))
builder.add_softmax_nd('classes_confidence', input_name='classes_raw:1', output_name='classes_confidence', axis=-1)

在上一篇文章中描述的 YOLO v2 预测解码中，我们为每个框返回一个最可能的类。Vision 框架希望我们为每个框返回 80 个类的置信度

builder.add_multiply_broadcastable(
    'combined_classes_confidence', 
    input_names=['box_confidence', 'classes_confidence'],
    output_name=‘combined_classes_confidence')

现在，我们拥有了我们需要的所有值。接下来，让我们将这些值格式化为 Vision 框架，将其转换为两个数组：一个包含所有边界框的坐标（每个框有四列），另一个包含为每个框/类组合计算的置信度（每个框有 80 列）。

这并不是一项困难的任务，但是由于我们需要将每个转换作为一个单独的操作来处理，因此它再次导致了冗余代码

builder.add_reshape_static('x:1', input_name='x', output_name='x:1', output_shape=(1,169*5,1))
builder.add_reshape_static('y:1', input_name='y', output_name='y:1', output_shape=(1,169*5,1))
builder.add_reshape_static('w:1', input_name='w', output_name='w:1', output_shape=(1,169*5,1))
builder.add_reshape_static('h:1', input_name='h', output_name='h:1', output_shape=(1,169*5,1))

builder.add_stack(
    'all_boxes:0', 
    input_names=['x:1', 'y:1', 'w:1', 'h:1'], 
    output_name='all_boxes:0', 
    axis=2)

builder.add_reshape_static(
    'all_boxes', 
    input_name='all_boxes:0', 
    output_name='all_boxes',
    output_shape=(1,169*5, 4))

builder.add_reshape_static(
    'all_scores', 
    input_name='combined_classes_confidence', 
    output_name='all_scores',
    output_shape=(1,169*5, 80))

使用格式化的 all_scores 和 all_boxes 数组，我们可以将这些数组映射到模型的输出并保存模型本身

builder.set_output(
    output_names= ['all_scores', 'all_boxes'],
    output_dims= [(845,80), (845,4)])

model_decoder = ct.models.MLModel(builder.spec)
model_decoder.save('./models/yolov2-decoder.mlmodel')

后续步骤

这涉及大量代码，但我们终于完成了。现在我们有了一个可以解码 YOLO v2 预测的 Core ML 模型。但是，如果没有与 YOLO 输出的链接，我们就无法使用它。在下一篇文章中，我们将创建一个 Core ML 管道，作为我们的端到端模型。

© . All rights reserved.