%use deeplearning4j-cuda(cuda=10.2)

//number of rows and columns in the input pictures
val numRows = 28
val numColumns = 28
val outputNum = 10 // number of output classes
val batchSize = 64 // batch size for each epoch
val rngSeed = 123 // random number seed for reproducibility
val numEpochs = 5 // number of epochs to perform
val rate = 0.0015 // learning rate

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator

//Get the DataSetIterators:
val mnistTrain = MnistDataSetIterator(batchSize, true, rngSeed)
val mnistTest = MnistDataSetIterator(batchSize, false, rngSeed)

import org.nd4j.linalg.learning.config.Nesterovs
import org.nd4j.linalg.lossfunctions.LossFunctions.LossFunction
import org.nd4j.linalg.activations.Activation

val conf = NeuralNetConfiguration.Builder()
        .seed(rngSeed.toLong()) //include a random seed for reproducibility
         // use stochastic gradient descent as an optimization algorithm

        .activation(Activation.RELU)
        .weightInit(WeightInit.XAVIER)
        .updater(Nesterovs(rate, 0.98)) //specify the rate of change of the learning rate.
        .l2(rate * 0.005) // regularize learning model
        .list()
        .layer(DenseLayer.Builder() //create the first input layer.
                .nIn(numRows * numColumns)
                .nOut(500)
                .build())
        .layer(DenseLayer.Builder() //create the second input layer
                .nIn(500)
                .nOut(100)
                .build())
        .layer(OutputLayer.Builder(LossFunction.NEGATIVELOGLIKELIHOOD) //create hidden layer
                .activation(Activation.SOFTMAX)
                .nIn(100)
                .nOut(outputNum)
                .build())
        .build()
        
val model = MultiLayerNetwork(conf)
model.init()
println(model.summary())

import org.deeplearning4j.ui.api.UIServer
import org.deeplearning4j.optimize.listeners.ScoreIterationListener
import org.deeplearning4j.api.storage.StatsStorageRouter
import org.deeplearning4j.api.storage.impl.RemoteUIStatsStorageRouter
import org.deeplearning4j.ui.stats.StatsListener

val uiServer: UIServer = UIServer.getInstance()
uiServer.enableRemoteListener()
//Create the remote stats storage router - this sends the results to the UI via HTTP, assuming the UI is at http://localhost:9000
val remoteUIRouter: StatsStorageRouter = RemoteUIStatsStorageRouter("http://localhost:9000")
model.setListeners(ScoreIterationListener(100), StatsListener(remoteUIRouter))

model.fit(mnistTrain, numEpochs)

val eval: org.nd4j.evaluation.classification.Evaluation = model.evaluate(mnistTest)
println(eval.stats())

uiServer.stop()