the img2col function in swift

main.swift
func img2col(_ input: [[[[Double]]]], kernelSize: Int, stride: Int) -> [[Double]] {
    let batchSize = input.count
    let inputHeight = input[0].count
    let inputWidth = input[0][0].count
    let inputChannel = input[0][0][0].count
    let outputHeight = (inputHeight - kernelSize) / stride + 1
    let outputWidth = (inputWidth - kernelSize) / stride + 1
    var output: [[Double]] = Array(repeating: Array(repeating: 0, count: kernelSize * kernelSize * inputChannel), count: batchSize * outputHeight * outputWidth)
    var outputIndex = 0
    for b in 0..<batchSize {
        for i in stride(from: 0, to: inputHeight - kernelSize + 1, by: stride) {
            for j in stride(from: 0, to: inputWidth - kernelSize + 1, by: stride) {
                for c in 0..<inputChannel {
                    for k in 0..<kernelSize {
                        for l in 0..<kernelSize {
                            output[outputIndex][c * kernelSize * kernelSize + k * kernelSize + l] = input[b][i + k][j + l][c]
                        }
                    }
                }
                outputIndex += 1
            }
        }
    }
    return output
}
1154 chars
26 lines