MLCompute_MNIST_Graph.swift

private func initializeTensors() {
    device = MLCDevice(type: .cpu)!

    inputTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [batchSize, imageSize, 1, 1], dataType: .float32)!)

    dense1WeightsTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [1, imageSize*dense1LayerOutputSize, 1, 1], dataType: .float32)!,
                                        randomInitializerType: .glorotUniform)
    dense1BiasesTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [1, dense1LayerOutputSize, 1, 1], dataType: .float32)!,
                                       randomInitializerType: .glorotUniform)
    dense2WeightsTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [1, dense1LayerOutputSize*numberOfClasses, 1, 1], dataType: .float32)!,
                                        randomInitializerType: .glorotUniform)
    dense2BiasesTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [1, numberOfClasses, 1, 1], dataType: .float32)!,
                                       randomInitializerType: .glorotUniform)

    lossLabelTensor = MLCTensor(descriptor: MLCTensorDescriptor(shape: [batchSize, numberOfClasses], dataType: .float32)!)
}

private func buildGraph() {
    graph = MLCGraph()

    dense1 = graph.node(with: MLCFullyConnectedLayer(weights: dense1WeightsTensor,
                                                        biases: dense1BiasesTensor,
                                                        descriptor: MLCConvolutionDescriptor(kernelSizes: (height: imageSize, width: dense1LayerOutputSize),
                                                                                             inputFeatureChannelCount: imageSize,
                                                                                             outputFeatureChannelCount: dense1LayerOutputSize))!,
                           sources: [inputTensor])

    relu1 = graph.node(with: MLCActivationLayer(descriptor: MLCActivationDescriptor(type: MLCActivationType.relu)!),
               source: dense1!)

    dense2 = graph.node(with: MLCFullyConnectedLayer(weights: dense2WeightsTensor,
                                                        biases: dense2BiasesTensor,
                                                        descriptor: MLCConvolutionDescriptor(kernelSizes: (height: dense1LayerOutputSize, width: numberOfClasses),
                                                                                             inputFeatureChannelCount: dense1LayerOutputSize,
                                                                                             outputFeatureChannelCount: numberOfClasses))!,
                           sources: [relu1!])

    outputSoftmax = graph.node(with: MLCSoftmaxLayer(operation: .softmax),
               source: dense2!)
}