对视频帧进行连续 Core ML 预测后出现内存泄漏

Question

我开发了一个应用程序，它从 WebRTC 接收视频帧，调整它们的大小，然后使用 Core ML 模型进行预测。源视频帧采用 420f PixelFormat 的 CVPixelBuffer 格式，我的 Core ML 模型不接受这种格式。因此，我使用 MTLTexture 来转换 CVPixelBuffer。这是相关代码：

// sourcePixelBuffer (420f) -> PixelBuffer (BGRA) & Resized PixelBuffer(BGRA)
public func processYUV420Frame2(sourcePixelBuffer: CVPixelBuffer, targetSize: CGSize)
 -> (original: CVPixelBuffer?, resized: CVPixelBuffer?) {
        guard let queue = self.commandQueue else {
            print("FrameMixer makeCommandBuffer failed")
            return (nil, nil)
        }
        
        let sourceWidth = CVPixelBufferGetWidth(sourcePixelBuffer)
        let sourceHeight = CVPixelBufferGetHeight(sourcePixelBuffer)
        
        guard let (yTexture, uvTexture) = createYUVTexturesFromPixelBuffer(sourcePixelBuffer) else {
            print("Failed to create YUV textures")
            return (nil, nil)
        }
        
        var originalBuffer: CVPixelBuffer?
        var resizedBuffer: CVPixelBuffer?
        
        autoreleasepool {
            let (originalTexture, resizedTexture) = convertYUVtoDualBGRA(
                device: metalDevice!,
                commandQueue: queue,
                yTexture: yTexture,
                uvTexture: uvTexture,
                sourceWidth: sourceWidth,
                sourceHeight: sourceHeight,
                targetWidth: Int(targetSize.width),
                targetHeight: Int(targetSize.height)
            )
            
            originalBuffer = createCVPixelBuffer(from: originalTexture)
            resizedBuffer = createCVPixelBuffer(from: resizedTexture)
        }
        
        return (originalBuffer, resizedBuffer)
    }

func createYUVTexturesFromPixelBuffer(_ pixelBuffer: CVPixelBuffer) -> (y: MTLTexture, uv: MTLTexture)? {
        guard let textureCache = textureCache else {
            print("make buffer failed, texture cache is not exist")
            return nil
        }
        
        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        
        var textureY: CVMetalTexture?
        var textureUV: CVMetalTexture?
        
        CVMetalTextureCacheFlush(textureCache, 0)
        
        // Create Y planer texture
        CVMetalTextureCacheCreateTextureFromImage(
            kCFAllocatorDefault,
            textureCache,
            pixelBuffer,
            nil,
            .r8Unorm,
            width,
            height,
            0,
            &textureY
        )
        
        // Create UV planer texture
        CVMetalTextureCacheCreateTextureFromImage(
            kCFAllocatorDefault,
            textureCache,
            pixelBuffer,
            nil,
            .rg8Unorm,
            width / 2,
            height / 2,
            1,
            &textureUV
        )
        
        guard let unwrappedTextureY = textureY, let unwrappedTextureUV = textureUV else {
            return nil
        }
        
        let y = CVMetalTextureGetTexture(unwrappedTextureY)!
        let uv = CVMetalTextureGetTexture(unwrappedTextureUV)!
        
        textureY = nil
        textureUV = nil
        return (y, uv)
    }

func convertYUVtoDualBGRA(device: MTLDevice, commandQueue: MTLCommandQueue, yTexture: MTLTexture, uvTexture: MTLTexture, sourceWidth: Int, sourceHeight: Int, targetWidth: Int, targetHeight: Int) -> (original: MTLTexture, resized: MTLTexture) {
        let originalTexture = createBGRATexture(device: device, width: sourceWidth, height: sourceHeight)
        let resizedTexture = createBGRATexture(device: device, width: targetWidth, height: targetHeight)
        
        guard let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else {
            fatalError("Failed to create command buffer or compute encoder")
        }
        
        computeEncoder.setComputePipelineState(dualOutputPipelineState!)
        computeEncoder.setTexture(yTexture, index: 0)
        computeEncoder.setTexture(uvTexture, index: 1)
        computeEncoder.setTexture(originalTexture, index: 2)
        computeEncoder.setTexture(resizedTexture, index: 3)
        
        var uniforms = DualOutputUniforms(sourceWidth: Float(sourceWidth), sourceHeight: Float(sourceHeight),
                                          targetWidth: Float(targetWidth), targetHeight: Float(targetHeight))
        computeEncoder.setBytes(&uniforms, length: MemoryLayout<DualOutputUniforms>.size, index: 0)
        
        let threadGroupSize = MTLSizeMake(16, 16, 1)
        let threadGroups = MTLSizeMake((sourceWidth + threadGroupSize.width - 1) / threadGroupSize.width,
                                       (sourceHeight + threadGroupSize.height - 1) / threadGroupSize.height,
                                       1)
        
        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
        computeEncoder.endEncoding()
        
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
        
        return (originalTexture, resizedTexture)
    }

此过程运行良好，没有内存问题，并且格式已正确从 420f 转换为 BGRA。但是，当使用调整大小的 CVPixelBuffer 作为 Core ML 模型的输入时，在连续处理视频帧时，资源无法正确释放。这可能会导致内存使用量高达 4.5GB，从而导致应用程序在 iPad 上崩溃。

有什么技巧或方法可以解决这个内存泄漏问题吗？使用 Core ML 连续处理视频帧时如何确保正确的资源管理？

Answer 1

我发现问题是下面的代码，我得到了返回的outputShapedArray来获取浮点数组：

    func runMidasWithResized(on pixelbuffer: CVPixelBuffer)
    -> [Float]?
    {
        var results: [Float]? = nil
        guard let prediction = try? mlmodel!.prediction(input: pixelbuffer) else {
            os_log("Prediction failed", type: .error)
            return nil
        }

        // Get result into format
        var shapedArray = prediction.outputShapedArray
        results = Array(repeating: 0.0, count: shapedArray.strides[0])
        shapedArray.withUnsafeMutableShapedBufferPointer { bufferPointer, shape, strides in
            results = Array(bufferPointer)
        }
        return results
    }
}

创建的结果似乎没有发布。

我修改了方法得到的结果如下：

        // Get result into format
        let output = prediction.output
        if let bufferPointer = try? UnsafeBufferPointer<Float>(output) {
            results = Array(bufferPointer)
        }
        return results

然后会正常释放创建的资源。