%install '.package(path: "$cwd/FastaiNotebook_04_callbacks")' FastaiNotebook_04_callbacks

import FastaiNotebook_04_callbacks
var data = mnistDataBunch(flat: true)

import TensorFlow

struct DynamicModel: Layer {
    public var hiddens: [Dense<Float>]
    public var last: Dense<Float>
    
    init(nIn: Int, nHiddens: [Int], nOut: Int) {
        var nLast = nIn
        self.hiddens = []
        for nHidden in nHiddens {
            self.hiddens.append(Dense.init(inputSize: nLast, outputSize: nHidden, activation: relu))
            nLast = nHidden
        }
        self.last = Dense.init(inputSize: nLast, outputSize: nOut, activation: identity)
    }
    
    @differentiable
    func applied(to input: Tensor<Float>, in context: Context) -> Tensor<Float> {
        return input.sequenced(in: context, through: hiddens, last)
    }
}

/// Model with a fixed number of layers, for comparison.
struct FixedModel: Layer {
    public var layer1: Dense<Float>
    public var layer2: Dense<Float>
    public var layer3: Dense<Float>
    
    public init(nIn: Int, nHid: Int, nOut: Int){
        layer1 = Dense(inputSize: nIn, outputSize: nHid, activation: relu)
        layer2 = Dense(inputSize: nHid, outputSize: nHid)
        layer3 = Dense(inputSize: nHid, outputSize: nOut)
    }
    
    @differentiable
    public func applied(to input: Tensor<Float>, in context: Context) -> Tensor<Float> {
        return input.sequenced(in: context, through: layer1, layer2, layer3)
    }
}

let nIn = 784
let nHid = 50
let nOut = 10

let modelDynamic = DynamicModel(nIn: nIn, nHiddens: [nHid, nHid], nOut: nOut)
var modelFixed = FixedModel(nIn: nIn, nHid: nHid, nOut: nOut)

// Turn off shuffling, and make the model weights the same, to verify that both models do the same thing.
data.shuffleTrain = false
modelFixed.layer1 = modelDynamic.hiddens[0]
modelFixed.layer2 = modelDynamic.hiddens[1]
modelFixed.layer3 = modelDynamic.last

let optDynamic = SimpleSGD<DynamicModel, Float>(learningRate: 1e-2)
func modelInitDynamic() -> DynamicModel { return modelDynamic }
let learnerDynamic = Learner(data: data, lossFunction: softmaxCrossEntropy, optimizer: optDynamic, initializingWith: modelInitDynamic)
learnerDynamic.delegates = [learnerDynamic.makeTrainEvalDelegate(), learnerDynamic.makeAvgMetric(metrics: [accuracy])]

let optFixed = SimpleSGD<FixedModel, Float>(learningRate: 1e-2)
func modelInitFixed() -> FixedModel { return modelFixed }
let learnerFixed = Learner(data: data, lossFunction: softmaxCrossEntropy, optimizer: optFixed, initializingWith: modelInitFixed)
learnerFixed.delegates = [learnerFixed.makeTrainEvalDelegate(), learnerFixed.makeAvgMetric(metrics: [accuracy])]

learnerDynamic.fit(3)

learnerFixed.fit(3)