Main.lean

import FunGrad

/-! Full-batch linear regression with shared slope and bias parameters. -/
namespace TrainingExample

def scalar := CoordinateSpace.scalar Float

def model := (BuildingBlocks.scale scalar).comp (BuildingBlocks.bias scalar)

/-- Half squared error, averaged across samples by the batching learner. -/
def squaredError : Backprop.ErrorData Float where
  derivative prediction target := prediction - target
  inverseDerivative input gradient := input - gradient

def learningRate : Float := 0.1

def sampleCount : Nat := 5

def inputs : Fin sampleCount → Float := ![-2.0, -1.0, 0.0, 1.0, 2.0]
def targets : Fin sampleCount → Float := ![-3.0, -1.0, 1.0, 3.0, 5.0]

/-- This works for any positive number of scalar input-target pairs. -/
def step (n : Nat) (hn : 0 < n) (xs ys : Fin n → Float)
    (parameters : Float × Float) : Float × Float :=
  (Batching.learner n hn learningRate squaredError model).U parameters xs ys

def meanLoss (parameters : Float × Float) : Float :=
  let errors := (List.finRange sampleCount).map fun i =>
    let prediction : Float := model.I parameters (inputs i)
    let error := prediction - targets i
    error * error / 2.0
  errors.foldr (· + ·) 0.0 / sampleCount.toFloat

def report (epoch : Nat) (parameters : Float × Float) : IO Unit :=
  IO.println s!"epoch {epoch}: slope={parameters.1}, bias={parameters.2}, mean loss={meanLoss parameters}"

end TrainingExample

def main : IO Unit := do
  IO.println "Full-batch linear regression: y = slope * x + bias."
  IO.println "Five samples from y = 2*x + 1; initial slope and bias are zero."
  IO.println s!"Learning rate: {TrainingExample.learningRate}; one averaged update per epoch."
  let mut parameters : Float × Float := (0.0, 0.0)
  TrainingExample.report 0 parameters
  for epoch in [1:101] do
    parameters := TrainingExample.step TrainingExample.sampleCount (by decide)
      TrainingExample.inputs TrainingExample.targets parameters
    if epoch % 10 == 0 then
      TrainingExample.report epoch parameters
  IO.println "Final predictions:"
  for i in List.finRange TrainingExample.sampleCount do
    let prediction : Float := TrainingExample.model.I parameters (TrainingExample.inputs i)
    IO.println s!"x={TrainingExample.inputs i}, target={TrainingExample.targets i}, prediction={prediction}"