coreml - CoreML Catalyst MPSNNReduceFeatureChannelsArgumentMax 返回全零

Question

我正在使用 CoreML 进行图像分割。我有一个 VNCoreMLRequest 来运行一个返回 MLMultiArray 的模型。为了加速处理模型输出，我使用 MPSNNReduceFeatureChannelsArgumentMax 将多阵列输出减少为 2D 阵列，然后将其转换为灰度图像。

这在 iOS 上效果很好，但在 Mac 上作为催化剂构建运行时，输出 2D 数组全为零。

我在 iMac Pro 上运行版本 12.2 beta 2 (12B5025f)。我没有看到任何运行时错误。MPSNNReduceFeatureChannelsArgumentMax 似乎不适用于 Mac Catalyst。我可以通过循环遍历所有数组维度来直接减少 cpu 上的通道，但这非常慢。这证明模型输出有效，只是金属减少功能失败。还有其他人使用 CoreML 和 Catalyst 吗？

这是不起作用的代码：

let buffer = self.queue.makeCommandBuffer()
let filter = MPSNNReduceFeatureChannelsArgumentMax(device: self.device)
filter.encode(commandBuffer: buffer!, sourceImage: probs, destinationImage: classes)

// add a callback to handle the buffer's completion and commit the buffer
buffer?.addCompletedHandler({ (_buffer) in
    let argmax = try! MLMultiArray(shape: [1, softmax.shape[1], softmax.shape[2]], dataType: .float32)
    classes.readBytes(argmax.dataPointer,
                      dataLayout: .featureChannelsxHeightxWidth,
                      imageIndex: 0)
    
    // unmap the discrete segmentation to RGB pixels
    guard var mask = codesToMask(argmax) else {
        return
    }

    // display image in view
    DispatchQueue.main.async {
        self.imageView.image = mask
    }
})

标签： coremlcatalystmac-catalyst