iOS中Core ML模型输出MLMultiArray的切片与重塑问题求助

嘿，作为Swift新手碰到Core ML的数组操作确实容易懵，我来一步步帮你搞定这个问题！Core ML的MLMultiArray本身没有内置的切片和重塑方法，不过我们可以通过直接操作内存指针来实现你的需求，而且效率也不错。

首先先明确几个前提：

• 你的finalOutput是**行优先（C-style）**的内存布局（Core ML默认就是这个布局），也就是维度顺序从外到内是[batch, height, width, channels]，内存中元素是按这个顺序连续存储的。
• 假设模型输出的数据类型是Float32（如果是Double或者其他类型，只需要替换对应的类型即可）。

完整实现步骤 #

1. 获取模型输出的MLMultiArray并验证

首先从模型输出中取出finalOutput，并确认它的形状符合预期：

// 假设modelOutput是你的Core ML模型的输出结果
guard let finalOutput = modelOutput.featureValue(for: "finalOutput")?.multiArrayValue else {
    print("Failed to retrieve finalOutput from model")
    return
}

// 验证形状是否正确
let expectedShape = [1, 40, 30, 90]
guard finalOutput.shape.map({ $0.intValue }) == expectedShape else {
    print("finalOutput shape mismatch: expected \(expectedShape), got \(finalOutput.shape.map({ $0.intValue }))")
    return
}

// 定义常量方便后续计算
let batchSize = 1
let height = 40
let width = 30
let totalChannels = 90
let elementSize = MemoryLayout<Float>.stride

// 获取数据指针，绑定到Float类型
guard let dataPointer = finalOutput.dataPointer.bindMemory(to: Float.self, capacity: finalOutput.count) else {
    print("Failed to bind memory to Float type")
    return
}

2. 处理第一个切片：(1,40,30,0:45) → 重塑为(1,10800,5)

这个切片对应每个(height, width)位置的前45个通道，我们需要将其重塑为10800个单元，每个单元5个特征：

// 创建目标MLMultiArray
let subArray1Shape = [NSNumber(value: batchSize), NSNumber(value: 10800), NSNumber(value: 5)]
guard var subArray1 = try? MLMultiArray(shape: subArray1Shape, dataType: .float32) else {
    print("Failed to create subArray1")
    return
}
let subArray1Pointer = subArray1.dataPointer.bindMemory(to: Float.self, capacity: subArray1.count)

// 高效复制数据：利用memcpy复制连续的小数据块
for h in 0..<height {
    for w in 0..<width {
        // 原数组中当前(height, width)块的起始偏移
        let baseSrcOffset = h * width * totalChannels + w * totalChannels
        // 目标数组中当前(height, width)对应的单元起始偏移
        let baseDstOffset = (h * width * 9 + w * 9) * 5
        
        for k in 0..<9 {
            // 每个k对应原数组中10个通道的前5个
            let srcStart = baseSrcOffset + k * 10
            let dstStart = baseDstOffset + k * 5
            // 复制5个连续元素
            memcpy(subArray1Pointer + dstStart, dataPointer + srcStart, 5 * elementSize)
        }
    }
}

3. 处理第二个切片：(1,40,30,45:54) → 重塑为(1,10800)

这个切片对应每个(height, width)位置的第45到53个通道（共9个），直接重塑为10800个单元素单元：

// 创建目标MLMultiArray
let subArray2Shape = [NSNumber(value: batchSize), NSNumber(value: 10800)]
guard var subArray2 = try? MLMultiArray(shape: subArray2Shape, dataType: .float32) else {
    print("Failed to create subArray2")
    return
}
let subArray2Pointer = subArray2.dataPointer.bindMemory(to: Float.self, capacity: subArray2.count)

// 高效复制：每个(height, width)块的9个通道是连续的，直接批量复制
for h in 0..<height {
    for w in 0..<width {
        let srcStart = h * width * totalChannels + w * totalChannels + 45
        let dstStart = h * width * 9 + w * 9
        memcpy(subArray2Pointer + dstStart, dataPointer + srcStart, 9 * elementSize)
    }
}

4. 处理第三个切片：(1,40,30,54:90) → 重塑为(1,10800,4)

这个切片对应每个(height, width)位置的第54到89个通道（共36个），重塑为10800个单元，每个单元4个特征：

// 创建目标MLMultiArray
let subArray3Shape = [NSNumber(value: batchSize), NSNumber(value: 10800), NSNumber(value: 4)]
guard var subArray3 = try? MLMultiArray(shape: subArray3Shape, dataType: .float32) else {
    print("Failed to create subArray3")
    return
}
let subArray3Pointer = subArray3.dataPointer.bindMemory(to: Float.self, capacity: subArray3.count)

// 高效复制数据
for h in 0..<height {
    for w in 0..<width {
        let baseSrcOffset = h * width * totalChannels + w * totalChannels + 54
        let baseDstOffset = (h * width * 9 + w * 9) * 4
        
        for k in 0..<9 {
            let srcStart = baseSrcOffset + k * 4
            let dstStart = baseDstOffset + k * 4
            memcpy(subArray3Pointer + dstStart, dataPointer + srcStart, 4 * elementSize)
        }
    }
}

关键注意事项 #

• 数据类型适配：如果你的模型输出是Double或者Int32，只需要把代码中的Float替换为对应的类型，同时修改MLMultiArray的dataType（比如.double、.int32）。
• 内存安全：确保在操作指针的过程中，finalOutput和新创建的MLMultiArray对象没有被提前释放，最好把这些操作放在同一个作用域内。
• 错误处理：示例中用了print和return处理错误，你可以根据自己的需求替换为更健壮的错误处理逻辑（比如抛出自定义错误）。

内容的提问来源于stack exchange，提问作者ScorpionMania