# + [markdown] colab_type="text" id="BcLhx7R4YuUv"
# # MNIST in Swift for TensorFlow (Perceptron)
#
# Blog post: https://rickwierenga.com/blog/s4tf/s4tf-mnist.html

# + [markdown] colab_type="text" id="lvZux3i-YzNv"
# ## Importing dependencies

# + colab={} colab_type="code" id="8GHBoSjRZuqt"
%install-location $cwd/swift-install
%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("tensorflow-0.6"))' Datasets

# + colab={} colab_type="code" id="dvaI1fPZRfQH"
import TensorFlow
import Foundation
import Datasets

# + colab={} colab_type="code" id="0awo42kpWWaE"
import Python
let plt = Python.import("matplotlib.pylab")
let np = Python.import("numpy")

# + colab={} colab_type="code" id="hxyQaWVajlrA"
%include "EnableIPythonDisplay.swift"
IPythonDisplay.shell.enable_matplotlib("inline")

# + [markdown] colab_type="text" id="gFp59EhiY3Y4"
# ## Loading MNIST
#
# [tensorflow/swift-models](https://github.com/tensorflow/swift-models)

# + colab={} colab_type="code" id="BBzeC88oZ1dS"
let batchSize = 512

# + colab={} colab_type="code" id="cmEkbL2UHP2r"
let mnist = MNIST(batchSize: batchSize, flattening: false, normalizing: true)

# + [markdown] colab_type="text" id="2ppaI_0SZKh2"
# ## Constructing the network

# + colab={} colab_type="code" id="k-hGkSXkRux3"
struct Model: Layer {
    var flatten = Flatten<Float>()
    var hiddenLayer = Dense<Float>(inputSize: 28 * 28, outputSize: 300, activation: relu)
    var outputLayer = Dense<Float>(inputSize: 300, outputSize: 10, activation: softmax)

    @differentiable
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        return input.sequenced(through: flatten, hiddenLayer, outputLayer)
    }
}

# + colab={} colab_type="code" id="5XCYHPH_qPaz"
var model = Model()

# + [markdown] colab_type="text" id="qAXVzcGFZTAn"
# ## Training

# + colab={} colab_type="code" id="px-gCWe7gSJp"
let epochs = 10
var trainHistory = np.zeros(epochs)
var valHistory = np.zeros(epochs)

# + colab={} colab_type="code" id="GvM422SwOyvN"
let optimizer = Adam(for: model)

# + colab={} colab_type="code" id="bdqKlxmpHkc7"
for epoch in 0..<epochs {
    // Update parameters
    Context.local.learningPhase = .training
    for i  in 0..<(mnist.trainingSize / batchSize)+1  {
        let thisBatchSize = i * batchSize >= mnist.trainingSize ? (mnist.trainingSize - ((i - 1) * batchSize)) : batchSize
        let images = mnist.trainingImages.minibatch(at: i, batchSize: thisBatchSize)
        let labels = mnist.trainingLabels.minibatch(at: i, batchSize: thisBatchSize)
        let (_, gradients) = valueWithGradient(at: model) { model -> Tensor<Float> in
            let logits = model(images)
            return softmaxCrossEntropy(logits: logits, labels: labels)
        }
        optimizer.update(&model, along: gradients)
    }

    // Evaluate model
    Context.local.learningPhase = .inference

    var correctTrainGuessCount = 0
    var totalTrainGuessCount = 0
    for i  in 0..<(mnist.trainingSize / batchSize)+1  {
        let thisBatchSize = i * batchSize >= mnist.trainingSize ? (mnist.trainingSize - ((i - 1) * batchSize)) : batchSize
        let images = mnist.trainingImages.minibatch(at: i, batchSize: thisBatchSize)
        let labels = mnist.trainingLabels.minibatch(at: i, batchSize: thisBatchSize)
        let logits = model(images)
        let correctPredictions = logits.argmax(squeezingAxis: 1) .== labels
        correctTrainGuessCount += Int(Tensor<Int32>(correctPredictions).sum().scalarized())
        totalTrainGuessCount += thisBatchSize
    }
    let trainAcc = Float(correctTrainGuessCount) / Float(totalTrainGuessCount)
    trainHistory[epoch] = PythonObject(trainAcc)

    var correctValGuessCount = 0
    var totalValGuessCount = 0
    for i  in 0..<(mnist.testSize / batchSize)+1  {
        let thisBatchSize = i * batchSize >= mnist.testSize ? (mnist.testSize - ((i - 1) * batchSize)) : batchSize
        let images = mnist.testImages.minibatch(at: i, batchSize: thisBatchSize)
        let labels = mnist.testLabels.minibatch(at: i, batchSize: thisBatchSize)
        let logits = model(images)
        let correctPredictions = logits.argmax(squeezingAxis: 1) .== labels
        correctValGuessCount += Int(Tensor<Int32>(correctPredictions).sum().scalarized())
        totalValGuessCount += thisBatchSize
    }
    let valAcc = Float(correctValGuessCount) / Float(totalValGuessCount)
    valHistory[epoch] = PythonObject(valAcc)
    
    print("\(epoch) | Training accuracy: \(trainAcc) | Validation accuracy: \(valAcc)")
}

# + [markdown] colab_type="text" id="WjUY0gwXTDX-"
# ## Inspecting training history

# + colab={} colab_type="code" id="b_f-DNeqp9qO"
plt.plot(trainHistory)
plt.title("Training History")
plt.show()

# + colab={} colab_type="code" id="1AkaTk_3THsK"
plt.plot(valHistory)
plt.title("Validation History")
plt.show()

StarcoderdataJupyterScriptsDedupFiltered