MPS Graph适配Denoiser至iOS时出现mps.matmul维度不匹配错误

正在将Facebook Denoiser模型适配到iOS设备，基于MPS Graph实现Conv1d与ConvTranspose1d的转2D适配逻辑时，出现以下错误：

(mpsFileLoc): /Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:39:0: error: 'mps.matmul' op contracting dimensions differ 1024 & 2048
(mpsFileLoc): /Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:39:0: note: see current operation: %212 = "mps.matmul"(%208#0, %211) <{transpose_lhs = false, transpose_rhs = false}> : (tensor<7x1x1024xf32>, tensor<2048x4096xf32>) -> tensor<7x1x4096xf32>

代码中未显式调用矩阵乘法，怀疑错误源于卷积/转置卷积的实现逻辑。以下是相关张量形状打印结果及核心代码：

卷积相关张量形状 #

1D->2D CONVTRANS input shape [1, 7, 1, 1024]
1D->2D CONVTRANS weight shape [1, 8, 1024, 512]
1D->2D CONVTRANS input shape [1, 32, 1, 512]
1D->2D CONVTRANS weight shape [1, 8, 512, 256]
1D->2D CONVTRANS input shape [1, 132, 1, 256]
1D->2D CONVTRANS weight shape [1, 8, 256, 128]
1D->2D CONVTRANS input shape [1, 532, 1, 128]
1D->2D CONVTRANS weight shape [1, 8, 128, 64]
1D->2D CONVTRANS input shape [1, 2132, 1, 64]
1D->2D CONVTRANS weight shape [1, 8, 64, 1]

ConvTranspose1D转2D实现代码 #

public func convolutionTranspose1D(source: MPSGraphTensor, weights: MPSGraphTensor, biases: MPSGraphTensor, outputShape: [NSNumber], descriptor: MPSGraphConvolution1dOpDescriptor, name: String? = nil) -> MPSGraphTensor {

    guard source.shape?.count == 3 else { fatalError("Source needs to be dimension of 4") }
    guard weights.shape?.count == 3 else { fatalError("Weights need to be dimension of 4") }
    var source = expandDims(source, axes: [2], name: name)
    source = transpose(source, permutation: [0, 3, 2, 1], name: name)
    print("1D->2D CONVTRANS input shape \(source.shape!)")
    var weights = expandDims(weights, axes: [2], name: name)
    weights = transpose(weights, permutation: [2, 3, 1, 0], name: name)
    print("1D->2D CONVTRANS weight shape \(weights.shape!)")
    var outputShape = outputShape
    outputShape.insert(1, at: 2)
    let conv2DDescriptor = descriptor.descriptor2D
    source = convolutionTranspose2D(source, weights: weights, outputShape: outputShape, descriptor: conv2DDescriptor, name: name)

    source = squeeze(source, axes: [2], name: name)

    return addition(source, biases, name: name)
}

private func convolutionTranspose1DLoad(source: MPSGraphTensor, chIn: Int, chOut: Int, padding: (Int, Int) = (0, 0), kernelSize: Int, stride: Int, dilation: Int = 1, path: String, name: String? = nil) -> MPSGraphTensor {

    let convWeights = fileLoad("\(path).weight", shape: [chOut as NSNumber, chIn as NSNumber, kernelSize as NSNumber], name: nil)
    let convBiases = fileLoad("\(path).bias", shape: [1, chOut as NSNumber, 1], name: nil)

    let (B, _, L) = source.shape!.dims3D
    let Lout = (L.intValue - 1) * stride - padding.0 - padding.1 + dilation * (kernelSize - 1) + 1
    
    let outputShape = [B, chOut as NSNumber, Lout as NSNumber]
    return convolutionTranspose1D(source: source, weights: convWeights, biases: convBiases, outputShape: outputShape, descriptor:  MPSGraphConvolution1dOpDescriptor(stride: stride, dilationRate: dilation, paddingLeft: padding.0, paddingRight: padding.1, paddingStyle: .explicit, dataLayout: .NCHW, weightsLayout: .HWIO))
}

权重加载与导出代码 #

Swift加载逻辑 #

private func fileLoad(_ file: String, shape: [NSNumber], name: String? = nil) -> MPSGraphTensor {
    let modelDataURL = Bundle.main.url(forResource: "\(file)", withExtension: "bin")
    let modelData = try! Data(contentsOf: modelDataURL!)
    return constant(modelData, shape: shape, dataType: .float32)
}

Python导出逻辑 #

for name, param in model.named_parameters():
    filename = str(name) + '.bin'
    path = os.path.join(basedir, filename)
    print(path)
    torch.save(param, path)

排查方向 #

1. 权重维度顺序错误：
   PyTorch中ConvTranspose1d的权重形状为(in_channels, out_channels, kernel_size)，但Swift加载时指定的形状是[chOut, chIn, kernelSize]，维度顺序完全相反，会导致通道维度错乱，进而引发内部matmul维度不匹配。
2. 转置逻辑不符合MPSGraph布局：
   MPSGraph ConvTranspose2D默认权重布局为HWIO（高度、宽度、输入通道、输出通道），当前转置后的权重形状与输入张量的布局是否匹配？若原始权重维度错误，转置后通道维度依然无法对齐。
3. 输出形状计算错误：
   ConvTranspose1D的输出长度公式是否正确？错误的输出形状会导致内部操作的维度不匹配。
4. 数据布局冲突：
   指定的dataLayout: .NCHW与转置后的输入张量形状是否一致？转置后的输入形状[1,7,1,1024]对应NCHW的话，空间维度与通道维度的顺序可能不符合1D卷积的预期，导致卷积操作时维度被错误解析。

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