连续Core ML视频帧预测后内存泄漏问题求助

问题背景 #

我开发了一款通过WebRTC接收视频帧、缩放后使用Core ML模型进行预测的应用。源视频帧为420f格式的CVPixelBuffer，无法被Core ML模型接受，因此使用MTLTexture进行格式转换（相关代码如下）。格式转换过程无内存问题，但将缩放后的CVPixelBuffer作为Core ML模型输入，连续处理视频帧时资源无法正常释放，内存占用最高达4.5GB，导致iPad端应用崩溃。

格式转换核心代码 #

// 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)
    }

内存泄漏解决与资源管理方案 #

1. 严格管理CVPixelBuffer生命周期 #

• 处理完Core ML推理后，手动调用CVPixelBufferRelease(resizedBuffer)释放输入帧，同时将引用置为nil，确保ARC能及时回收内存。
• 检查createCVPixelBuffer(from:)方法的实现，确保使用kCFAllocatorDefault或自定义可回收分配器，避免创建无法自动释放的缓冲区。
• 将Core ML推理逻辑包裹在autoreleasepool中，加速临时对象的释放：

  autoreleasepool {
      // 执行Core ML推理
      let output = try model.prediction(input: input)
      // 处理输出
      inputBuffer = nil // 立即置空输入缓冲区引用
  }
  

2. Core ML模型与推理优化 #

• 复用模型实例：全局维护一个Core ML模型实例，不要每次处理帧都重新初始化，避免重复分配模型权重内存。
• 使用Vision框架封装推理：Vision的VNCoreMLRequest会自动管理内存和资源复用，比直接调用Core ML更适配连续帧场景，示例代码：

  // 全局初始化一次
  let config = MLModelConfiguration()
  config.computeUnits = .neuralEngine // 优先使用神经引擎，减少内存占用
  let coreMLModel = try YourCoreMLModel(configuration: config).model
  let vnModel = try VNCoreMLModel(for: coreMLModel)
  let request = VNCoreMLRequest(model: vnModel) { [weak self] request, error in
      guard let results = request.results as? [VNClassificationObservation], error == nil else {
          return
      }
      // 处理推理结果
      self?.currentResizedBuffer = nil // 释放当前帧引用
  }
  
  // 每帧处理
  guard let buffer = resizedBuffer else { return }
  let handler = VNImageRequestHandler(cvPixelBuffer: buffer, options: [:])
  try handler.perform([request])
  

• 配置模型内存限制：通过MLModelConfiguration启用低精度计算减少内存开销：

  config.allowsLowPrecisionAccumulation = true
  

3. Metal资源的额外优化 #

• 避免在convertYUVtoDualBGRA中使用commandBuffer.waitUntilCompleted()同步等待，改用异步回调减少线程阻塞和资源堆积：

  commandBuffer.addCompletedHandler { [weak self] _ in
      autoreleasepool {
          self?.originalBuffer = createCVPixelBuffer(from: originalTexture)
          self?.resizedBuffer = createCVPixelBuffer(from: resizedTexture)
          // 触发后续Core ML推理
          self?.processCoreMLInference()
      }
  }
  commandBuffer.commit()
  

• 优化CVMetalTextureCache使用：不要每次创建纹理都调用CVMetalTextureCacheFlush(textureCache, 0)，改为定期清理（比如每处理100帧），避免缓存频繁重建带来的内存波动。

4. 帧处理节流与队列管控 #

• 如果WebRTC帧率高于Core ML处理能力，实现帧丢弃逻辑：只处理最新的帧，丢弃堆积的旧帧，避免内存中积压大量未处理的CVPixelBuffer。
• 使用串行队列处理Core ML推理，确保同一时间只有一个推理任务执行，避免多线程下的资源竞争和内存暴涨：

  let inferenceQueue = DispatchQueue(label: "com.yourapp.coreml.inference")
  inferenceQueue.async {
      // 执行推理逻辑
  }
  

5. 泄漏排查工具使用 #

• 用Xcode的Memory Graph Debugger查看对象引用链，定位持有CVPixelBuffer或Core ML对象的强引用（比如闭包未使用weak self导致的循环引用）。
• 使用Instruments的Leaks和Allocations工具，跟踪内存分配热点，找到未释放的大内存对象。

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