MPSCNN 权重排序

Question

Metal Performance Shader 框架支持构建您自己的卷积神经网络。例如，当创建一个 MSPCNNConvolution 时，它需要一个 4D 权重张量作为 init 参数，表示为一维浮点指针。

init(device: MTLDevice,
  convolutionDescriptor: MPSCNNConvolutionDescriptor,
  kernelWeights: UnsafePointer<Float>,
  biasTerms: UnsafePointer<Float>?,
  flags: MPSCNNConvolutionFlags)

文档中有关于 4D 张量的内容

The layout of the filter weight is arranged so that it can be reinterpreted as a 4D tensor (array) weight[outputChannels][kernelHeight][kernelWidth][inputChannels/groups]

不幸的是，这些信息并没有真正告诉我如何将 4D 数组排列成一维 Float 指针。

我尝试按照 BNNS 对方要求的方式订购权重，但没有成功。

如何将 4D 张量（数组）正确表示为 1D Float 指针（数组）？

PS: 我试过像 C 数组一样排列它并获取指向平面数组的指针，但它没有用。

更新

@RhythmicFistman：这就是我将它存储在普通数组中的方式，我可以将其转换为 UsafePointer<Float>（但不起作用）：

var output = Array<Float>(repeating: 0, count: weights.count)

for o in 0..<outputChannels {
    for ky in 0..<kernelHeight {
        for kx in 0..<kernelWidth {
            for i in 0..<inputChannels {
                let offset = ((o * kernelHeight + ky) * kernelWidth + kx) * inputChannels + i
                output[offset] = ...
            }
        }
    }
}

Answer 1

好的，我明白了。这是我用来改造卷积和全连接矩阵的 2 python 函数

# shape required for MPSCNN [oC kH kW iC]
# tensorflow order is [kH kW iC oC]
def convshape(a):
    a = np.swapaxes(a, 2, 3)
    a = np.swapaxes(a, 1, 2)
    a = np.swapaxes(a, 0, 1)
    return a

# fully connected only requires a x/y swap
def fullshape(a):
    a = np.swapaxes(a, 0, 1)
    return a

Answer 2

这是我最近必须为 Caffe 权重做的事情，因此我可以提供 Swift 实现来说明我如何重新排序这些权重。下面的函数接收一个用于卷积的 Caffe 权重的 Float 数组（按 [c_o][c_i][h][w] 顺序）并将它们重新排序为 Metal 所期望的（[c_o][h][w][c_i]顺序):

public func convertCaffeWeightsToMPS(_ weights:[Float], kernelSize:(width:Int, height:Int), inputChannels:Int, outputChannels:Int, groups:Int) -> [Float] {

    var weightArray:[Float] = Array(repeating:0.0, count:weights.count)
    var outputIndex = 0

    let groupedInputChannels = inputChannels / groups
    let outputChannelWidth = groupedInputChannels * kernelSize.width * kernelSize.height

    // MPS ordering: [c_o][h][w][c_i]
    for outputChannel in 0..<outputChannels {
        for heightInKernel in 0..<kernelSize.height {
            for widthInKernel in 0..<kernelSize.width {
                for inputChannel in 0..<groupedInputChannels {
                    // Caffe ordering: [c_o][c_i][h][w]
                    let calculatedIndex = outputChannel * outputChannelWidth + inputChannel * kernelSize.width * kernelSize.height + heightInKernel * kernelSize.width + widthInKernel
                    weightArray[outputIndex] = weights[calculatedIndex]
                    outputIndex += 1
                }
            }
        }
    }

    return weightArray
}

根据我的图层可视化，这似乎生成了正确的卷积结果（与 Caffe 生成的结果相匹配）。我相信它也适当地考虑了分组，但我需要验证这一点。

Tensorflow 的顺序与 Caffe 不同，但您应该能够更改循环内部的数学来解决这个问题。

Answer 3

此处的文档假定您具有一些 C 语言的专业知识。在这种情况下，当 x、y 和 z 是编译时已知的常量时，a[x][y][z] 通常会折叠成一维数组。发生这种情况时，z 分量变化最快，其次是 y，然后是 x -- 从外到内。

如果我们有 a[2][2][2]，它会折叠成 1D 为：

{ a[0][0][0], a[0][0][1], a[0][1][0], a[0][1][1], 
  a[1][0][0], a[1][0][1], a[1][1][0], a[1][1][1] }

Answer 4

我认为tensorflow已经有一个方便的方法来完成这样的任务：

tf.transpose(aWeightTensor, perm=[3, 0, 1, 2])

完整文档：https://www.tensorflow.org/api_docs/python/tf/transpose

MPSCNN 权重排序

MPSCNN Weight Ordering

neural-network

ios

swift

metal

conv-neural-network