MAuto.swift

import Foundation
import TensorFlow

struct MAuto: Parameterized {

    // Define constants
    static let imageSize: Int32 = 28
    static let inputSize: Int32 = imageSize * imageSize
    static let hiddenLength: Int32 = 512
    static let coderLength: Int32 = 256

    // Declare weights as TF Parameters
    @TFParameter var w1: Tensor<Float> // 784 -> 512
    @TFParameter var w2: Tensor<Float> // 512 -> 256
    @TFParameter var w3: Tensor<Float> // 256 -> 512
    @TFParameter var w4: Tensor<Float> // 512 -> 784

    // Define learning rate
    var learningRate: Float = 0.001

    init() {
        let w1 = Tensor<Float>(randomUniform: [MAuto.inputSize, MAuto.hiddenLength])
        let w2 = Tensor<Float>(randomUniform: [MAuto.hiddenLength, MAuto.coderLength])
        let w3 = Tensor<Float>(randomUniform: [MAuto.coderLength, MAuto.hiddenLength])
        let w4 = Tensor<Float>(randomUniform: [MAuto.hiddenLength, MAuto.inputSize])

        // Xavier weight initialization method
        self.w1 = w1 / sqrtf(Float(MAuto.inputSize))
        self.w2 = w2 / sqrtf(Float(MAuto.hiddenLength))
        self.w3 = w3 / sqrtf(Float(MAuto.coderLength))
        self.w4 = w4 / sqrtf(Float(MAuto.hiddenLength))
    }

    // Returns embedding for input (layer 2 output)
    func embedding(input: Tensor<Float>) -> Tensor<Float> {
        let o1 = tanh(matmul(input, w1))
        let o2 = tanh(matmul(o1, w2))
        return o2
    }

    // Returns decoding (image) for embedding
    func decodingFor(embedding: Tensor<Float>) -> Tensor<Float> {
        let o3 = tanh(matmul(embedding, w3))
        let o4 = sigmoid(matmul(o3, w4))
        return o4
    }

    // Predict and get loss
    func predict(input: Tensor<Float>) -> (loss: Float, embedding: Tensor<Float>, decoding: Tensor<Float>) {
        let embedding = self.embedding(input: input)
        let decoding = self.decodingFor(embedding: embedding)
        let loss: Float = 0.5 * (decoding - input).squared().mean()
        return (loss: loss, embedding: embedding, decoding: decoding)
    }

    // Trains 1 step
    mutating func trainStep(input: Tensor<Float>) -> Float {
        let learningRate = self.learningRate
        let batchSize = Tensor<Float>(input.shapeTensor[0])

        let o1 = matmul(input, w1)
        let onl1 = tanh(o1)
        
        let o2 = matmul(onl1, w2)
        let onl2 = tanh(o2)

        let o3 = matmul(onl2, w3)
        let onl3 = tanh(o3)

        let o4 = matmul(onl3, w4)
        let onl4 = sigmoid(o4)
        let predictions = onl4

        let dz4 = (predictions - input) / batchSize
        let dw4 = matmul(o3.transposed(), dz4)
        
        let dz3 = matmul(dz4, w4.transposed()) * (1 - o3.squared())
        let dw3 = matmul(o2.transposed(), dz3)

        let dz2 = matmul(dz3, w3.transposed()) * (1 - o2.squared())
        let dw2 = matmul(o1.transposed(), dz2)

        let dz1 = matmul(dz2, w2.transposed()) * (1 - o1.squared())
        let dw1 = matmul(input.transposed(), dz1)

        let gradients = Parameters(w1: dw1, w2: dw2, w3: dw3, w4: dw4)

        let loss = 0.5 * (predictions - input).squared().mean()

        allParameters.update(withGradients: gradients) { p, g in
            p -= g * learningRate
        }

        return loss
    }

    mutating func train(images: Tensor<Float>, iterationCount: Int32) {
        let batchSize: Int32 = 64
        for i in 1...iterationCount {
            for batchStep in 0..<930 {
                let batch = batchSize * Int32(batchStep)
                let images = images.slice(lowerBounds: [batch, 0], upperBounds: [batch + batchSize, MAuto.inputSize])
                let loss = trainStep(input: images)
                if i % 5 == 0 && batchStep == 0 {
                    print("\(i) steps, loss: \(loss)")
                }
            }
        }
    }

}

let autoencoder = MAuto()