MLCompute_MNIST_TrainingLoop.swift

private func execTrainingLoop(log: (String) -> Void) {
    let trainingSample = trainingDataX!.count / imageSize
    let trainingBatches = trainingSample / batchSize

    for epoch in 0..<epochs {
        var epochMatch = 0

        for batch in 0..<trainingBatches {
            let xData = trainingDataX!.withUnsafeBufferPointer { pointer in
                MLCTensorData(immutableBytesNoCopy: pointer.baseAddress!.advanced(by: batch * imageSize * batchSize),
                              length: batchSize * imageSize * MemoryLayout<Float>.size)
            }

            let yData = trainingDataY!.withUnsafeBufferPointer { pointer in
                MLCTensorData(immutableBytesNoCopy: pointer.baseAddress!.advanced(by: batch * numberOfClasses * batchSize),
                              length: batchSize * numberOfClasses * MemoryLayout<Int>.size)
            }

            trainingGraph.execute(inputsData: ["image" : xData],
                                  lossLabelsData: ["label" : yData],
                                  lossLabelWeightsData: nil,
                                  batchSize: batchSize,
                                  options: [.synchronous]) { [self] (r, e, time) in
                // VALIDATE
                let bufferOutput = UnsafeMutableRawPointer.allocate(byteCount: batchSize * self.numberOfClasses * MemoryLayout<Float>.size, alignment: MemoryLayout<Float>.alignment)

                outputSoftmax!.copyDataFromDeviceMemory(toBytes: bufferOutput, length: batchSize * self.numberOfClasses * MemoryLayout<Float>.size, synchronizeWithDevice: false)

                let float4Ptr = bufferOutput.bindMemory(to: Float.self, capacity: batchSize * self.numberOfClasses)
                let float4Buffer = UnsafeBufferPointer(start: float4Ptr, count: batchSize * self.numberOfClasses)
                let batchOutputArray = Array(float4Buffer)

                for i in 0..<batchSize {
                    let batchStartingPoint = i * self.numberOfClasses
                    let predictionStartingPoint = (i * self.numberOfClasses) + (batch * batchSize * numberOfClasses)
                    let sampleOutputArray = Array(batchOutputArray[batchStartingPoint..<(batchStartingPoint + self.numberOfClasses)])
                    let predictionArray = Array(trainingDataY![predictionStartingPoint..<(predictionStartingPoint + numberOfClasses)])

                    let prediction = argmaxDecoding(sampleOutputArray)
                    let label = oneHotDecoding(predictionArray)

                    if prediction == label {
                        epochMatch += 1
                    }
                }
            }
        }

        let epochAccuracy = Float(epochMatch) / Float(trainingSample)
        log("Epoch \(epoch) Accuracy = \(epochAccuracy) %")
    }
}