iOS（Swift4）基于MPS复现YOLOv2：space_to_depth与卷积拼接实现咨询

Hey there! I’ve tackled similar MPS + YOLOv2 challenges on iOS (Swift 4) before, so let’s break down practical solutions for both your pain points:

MPS doesn’t have a native kernel for this operation, but we can build a custom MPSUnaryImageKernel to handle it. For YOLOv2, we use a block size of 2—this takes a 26x26x256 feature map and converts it to 13x13x1024 by moving spatial blocks into the channel dimension.

Step 1: Custom Swift Kernel Class #

This class wraps our Metal shader and validates input/output dimensions to avoid runtime errors:

import MetalPerformanceShaders

class MPSImageSpaceToDepth: MPSUnaryImageKernel {
    let blockSize: UInt
    
    init(device: MTLDevice, blockSize: UInt = 2) {
        self.blockSize = blockSize
        super.init(device: device)
        self.label = "MPSImageSpaceToDepth_YOLOv2"
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override func encode(commandBuffer: MTLCommandBuffer, sourceImage: MPSImage, destinationImage: MPSImage) {
        // Enforce valid dimensions for YOLOv2's block size 2
        precondition(sourceImage.height == destinationImage.height * Int(blockSize), "Input height must be blockSize × output height")
        precondition(sourceImage.width == destinationImage.width * Int(blockSize), "Input width must be blockSize × output width")
        precondition(destinationImage.featureChannels == sourceImage.featureChannels * Int(blockSize * blockSize), "Output channels = input channels × blockSize²")
        super.encode(commandBuffer: commandBuffer, sourceImage: sourceImage, destinationImage: destinationImage)
    }
    
    override func newDevice(_ device: MTLDevice) -> MPSUnaryImageKernel {
        return MPSImageSpaceToDepth(device: device, blockSize: blockSize)
    }
    
    override var kernelFunctionName: String {
        return "space_to_depth_yolov2_kernel"
    }
}

Step 2: Metal Shader Implementation #

Add this to your .metal file—it handles the core spatial-to-channel mapping logic:

#include <metal_stdlib>
#include <MetalPerformanceShaders/MetalPerformanceShaders.h>

using namespace metal;

kernel void space_to_depth_yolov2_kernel(
    texture2d_array<float, access::read>  inTexture  [[texture(0)]],
    texture2d_array<float, access::write> outTexture [[texture(1)]],
    uint2 gridPosition                    [[thread_position_in_grid]],
    uint blockSize                        [[constant(0)]]
) {
    uint outX = gridPosition.x;
    uint outY = gridPosition.y;
    uint inChannels = inTexture.get_array_size();
    
    // Map output coordinates to input's block starting point
    uint inX = outX * blockSize;
    uint inY = outY * blockSize;
    
    // Iterate over block pixels and map to channel dimension
    for (uint by = 0; by < blockSize; by++) {
        for (uint bx = 0; bx < blockSize; bx++) {
            uint inTexX = inX + bx;
            uint inTexY = inY + by;
            // Calculate channel offset for the current block position
            uint channelOffset = (by * blockSize + bx) * inChannels;
            
            for (uint c = 0; c < inChannels; c++) {
                float value = inTexture.read(uint3(inTexX, inTexY, c)).x;
                outTexture.write(float4(value, 0.0, 0.0, 0.0), uint3(outX, outY, channelOffset + c));
            }
        }
    }
}

Step 3: Integrate into Your YOLOv2 Pipeline #

// Assuming conv18 outputs a 26x26x256 MPSImage
let device = MTLCreateSystemDefaultDevice()!
let commandBuffer = device.makeCommandQueue()!.makeCommandBuffer()!

// Create output image for space-to-depth (13x13x1024)
let stdOutputDescriptor = MPSImageDescriptor(
    width: 13,
    height: 13,
    featureChannels: 256 * 4, // 256 * 2²
    pixelFormat: .float16 // Match your input pixel format
)
let stdResult = MPSImage(device: device, imageDescriptor: stdOutputDescriptor)

// Run the custom kernel
let spaceToDepth = MPSImageSpaceToDepth(device: device)
spaceToDepth.encode(commandBuffer: commandBuffer, sourceImage: conv18.resultImage, destinationImage: stdResult)

Luckily, MPS has a native kernel for this exact use case: MPSImageAppendChannels. It concatenates feature maps along the channel dimension, as long as all inputs share identical width, height, and pixel format.

Implementation Code #

// Assuming you have your two input feature maps:
// stdResult = 13x13x256 (from space-to-depth)
// conv19Result = 13x13x1024 (from your existing conv19 node)
let conv19Result = conv19.resultImage

// Create output descriptor for concatenated feature map
let concatDescriptor = MPSImageDescriptor(
    width: 13,
    height: 13,
    featureChannels: 256 + 1024, // 1280 total channels
    pixelFormat: conv19Result.pixelFormat // Match input format
)
let concatResult = MPSImage(device: device, imageDescriptor: concatDescriptor)

// Initialize and run the concatenation kernel
let appendChannels = MPSImageAppendChannels(device: device)
// The order of sourceImages defines the channel order in the output
appendChannels.encode(commandBuffer: commandBuffer, sourceImages: [stdResult, conv19Result], destinationImage: concatResult)

// Don't forget to commit the command buffer!
commandBuffer.commit()

Key Notes #

• Ensure all input images use the same pixel format (e.g., .float16 for memory efficiency on iOS)
• Double-check feature channel counts to avoid runtime crashes
• For YOLOv2, the concatenation order matters—verify you’re appending the 256-channel feature map from the earlier layer in the correct position relative to the 1024-channel map.

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