MPSCNN权重排序

Question

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

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

关于4D张量，文献是这么说的

过滤器权重的布置使其可以重新解释为4D张量(数组)体重[outputChannels] [kernelHeight] [kernelWidth] [inputChannels/组]

不幸的是，这些信息并没有真正告诉我如何将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

好了，我算出来了。下面是我用来改造卷积和全连通矩阵的两个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权值进行卷积(以[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——从外到内。

如果我们有一个[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