使用数组运算解码 YOLO Core ML 对象检测器

• 下载解码向量化 - 1.3 MB

引言

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

正确解码 YOLO 输出

如果您之前使用过神经网络或数组，您很可能会在看到我们之前文章中的单元格和框的循环（cy, cx 和 b）时感到不舒服。通常，如果您在使用数组时需要循环，那么您就做错了。在这种特殊情况下，这是故意的，因为这些循环使理解底层逻辑更容易。向量化实现通常很短，但乍一看并不容易理解。

请注意，本文代码下载中的笔记本包含了以前的（基于循环的）解决方案和新的解决方案。

要开始向量化解码，我们需要一个新的 softmax 函数，该函数在二维数组上运行

def softmax_2d(x, axis=1):
    x_max = np.max(x, axis=axis)[:, np.newaxis]
    e_x = np.exp(x - x_max)
    x_sum = np.sum(e_x, axis=axis)[:, np.newaxis]    
    return e_x / x_sum

接下来，为了摆脱 cy、cx 和 b 循环，我们需要一些常量数组

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, GRID_SIZE**2, 1)
CY = np.tile(np.arange(GRID_SIZE), GRID_SIZE).reshape(1, GRID_SIZE, GRID_SIZE).transpose()
CY = CY.reshape(1, GRID_SIZE**2, 1)

ANCHORS 数组现在被分成两个：ANCHORS_W 和 ANCHORS_H。

CX 和 CY 数组包含所有 cx 和 cy 的值组合，这些值以前是在嵌套循环执行期间生成的。这些数组的形状设置为简化后续操作。

现在我们准备实现向量化解码函数

def decode_preds_vec(raw_preds: []):
    num_classes = len(COCO_CLASSES)

    raw_preds = np.transpose(raw_preds, (0, 2, 3, 1))

    raw_preds = raw_preds.reshape((1, GRID_SIZE**2, BOXES_PER_CELL, num_classes + 5))
    decoded_preds = []
    
    tx = raw_preds[:,:,:,0]
    ty = raw_preds[:,:,:,1]
    tw = raw_preds[:,:,:,2]
    th = raw_preds[:,:,:,3]
    tc = raw_preds[:,:,:,4]
    
    x = ((CX + sigmoid(tx)) * CELL_SIZE).reshape(-1)
    y = ((CY + sigmoid(ty)) * CELL_SIZE).reshape(-1)
    
    w = (np.exp(tw) * ANCHORS_W * CELL_SIZE).reshape(-1)
    h = (np.exp(th) * ANCHORS_H * CELL_SIZE).reshape(-1)
    
    box_confidence = sigmoid(tc).reshape(-1)
    
    classes_raw = raw_preds[:,:,:,5:5 + num_classes].reshape(GRID_SIZE**2 * BOXES_PER_CELL, -1)
    classes_confidence = softmax_2d(classes_raw, axis=1)
    
    box_class_idx = np.argmax(classes_confidence, axis=1)   
    box_class_confidence = classes_confidence.max(axis=1)
    combined_box_confidence = box_confidence * box_class_confidence
    
    decoded_boxes = np.stack([
        box_class_idx, 
        combined_box_confidence, 
        x, 
        y, 
        w, 
        h]).transpose()
    
    return sorted(list(decoded_boxes), key=lambda p: p[1], reverse=True)

首先，为了使计算稍微容易一些，我们通过将带有编码的框坐标和类置信度的 425 个值移动到最后一个维度来转置 raw_preds 数组。然后我们将其从 (1, 13, 13, 425) 调整为 (1, 13*13, 5, 85)。这样，忽略第一个位置的批次（始终等于 0），维度的顺序与之前的 cy (13)、cx (13) 和 box (5) 循环匹配。

请注意，我们必须使用形状 (1, 13*13, 5, 85)，而不是更明确的 (1, 13, 13, 5, 85)，这仅仅是因为 Core ML 存在一些数组秩限制。这意味着某些操作会导致对具有太多维度的数组的异常。此外，考虑到“隐藏”的内部序列维度，在 Core ML 中使用数组并不是很直观。

在 NumPy 数组上工作时，我们可以使用“更长”的形状 (1, 13, 13, 5, 85)，但是，为了使操作易于转换为 Core ML，我们必须将维度数量减少一，因此形状为 (1, 13*13, 5, 85)。

现在，与之前的版本相比，主要的变化在于如何获取 tx、ty、tw、th 和 tc 值以及 classes_raw。我们不是读取对应于单个单元格内单个框的单独值，而是在一个步骤中获取一个包含所有对应值的数组。这支持以下“单步”数组运算，这使得所有计算都非常高效，尤其是在为数组计算优化的芯片（例如 GPU 或神经引擎）上执行时。

decoded_preds_vec = decode_preds_vec(preds)
annotate_image(image, decoded_preds_vec)

这里是另一个例子。

后续步骤

现在我们得到了与以前的基于循环的解决方案相同的结果。这使我们准备好将检测解码直接包含在 Core ML 模型中。这最终将允许我们使用 Vision 框架 的对象检测功能，这大大简化了 iOS 应用程序的 Swift 代码。