Implement training #31

// Define the model architecture
val model = network {
    input(784) // 28x28 = 784 pixels
    dense(128) {
        activation = ReLU()::forward
    }
    dense(64) {
        activation = ReLU()::forward
    }
    dense(10) // 10 output classes (digits 0-9)
}

// Configure training
val trainer = trainer {
    model(model)
    loss(CrossEntropyLoss())
    optimizer(Adam()) {
        learningRate = 0.001
        weightDecay = 1e-5
    }
    batchSize = 64
    epochs = 10
    
    validation {
        splitRatio = 0.1
        metric = Accuracy()
        frequency = 1 // validate every epoch
    }
    
    callbacks {
        onEpochEnd {
            println("Epoch ${it.epoch}: loss=${it.loss}, accuracy=${it.accuracy}")
        }
        earlyStop {
            patience = 3
            metric = "val_loss"
            mode = "min"
        }
    }
}

// Train the model
trainer.fit(trainImages, trainLabels)

// Evaluate the model
val evaluation = trainer.evaluate(testImages, testLabels)
println("Test accuracy: ${evaluation.accuracy}")

// Save the trained model
model.save("mnist_model.gguf")