0 | module NN.Training.Training
  1 |
  2 | import Data.Tensor
  3 | import Data.Container.Additive as Additive
  4 | import public Data.ScientificNotation
  5 | import NN.Optimisers
  6 |
  7 | import NN.Utils
  8 | import NN.Training.DataLoader
  9 | import Data.Para
 10 | import Data.Autodiff.Model
 11 |
 12 | {-------------------------------------------------------------------------------
 13 | {-------------------------------------------------------------------------------
 14 | TODO update this in light of effects
 15 |
 16 | This file defines functions which perform pure optimisation:
 17 | optimisation of a differentiable function `f : p -> x` (here modelled as a lens `f : p =%> x` via some optimiser such as gradient descent. 
 18 | Here no "loss function" or "input-output pairs" are needed, just a function to optimise.
 19 |
 20 | This file provides functionality for creating, and turning supervised learning problems into pure optimisation problems, via a function which takes:
 21 | a) a parametric lens `f : x >< p =%>> y`
 22 | b) a loss function `loss : (y, y) =%> l`
 23 | c) input-output pairs `IO (x.Shp, y.Shp)`
 24 | and composes them to produce an optimisation problem `f : p =%> l` which the above described functions can consume.
 25 |
 26 | Notably, only a *non-dependent* supervised-learning problem can be turned into a pure optimisation one. If the parameter space depends on the input, then learning becomes its own thing.
 27 |
 28 |
 29 | todo using Hom-version of optimisation becomes problematic if we either
 30 | a) have the dependency of the parameter on the input
 31 | b) use monadic lenses
 32 |
 33 | -------------------------------------------------------------------------------}
 34 | -------------------------------------------------------------------------------}
 35 |
 36 | ||| Performs a single step of optimisation of some differentiable function
 37 | ||| `f : p -> l`, additionally handling some effect `e`
 38 | ||| The optimiser used is allowed to be stateful meaning the result of the
 39 | ||| optimisation is both the final parameter and the state of the optimiser
 40 | public export
 41 | optimiseStep : {p, l : AddCont} -> {e : Cont} ->
 42 |   InterfaceOnPositions l Num =>
 43 |   (f : p =%+> e >-+@ l) ->
 44 |   (handleEffect : Costate (IO <!> e)) ->
 45 |   (optimiser : Optimiser p stateTy) ->
 46 |   Costate (IO <!> (Const (p.Shp, stateTy)))
 47 | optimiseStep f handleEffect (MkOptimiser opt _) =
 48 |   let closeFunction : p =%+> !* e
 49 |       closeFunction = f %+>> (id >-+@ constantOne) %+>> actionToFree
 50 |
 51 |       closeFunctionT : UC p =%> e -- transposed variant of closeFunction
 52 |       closeFunctionT = addContTransposeInv closeFunction
 53 |
 54 |   in (IO <!> (opt %>> closeFunctionT)) %>> handleEffect
 55 |
 56 | ||| Evaluates the forward pass of some effectful lens
 57 | public export
 58 | evalFw : {0 e : Cont} ->
 59 |   (f : a -> Ext e b) ->
 60 |   (handleEffect : Costate (IO <!> e)) ->
 61 |   Costate (IO <!> (Const2 a b))
 62 | evalFw f handleEffect = toCostate $ \ps => do
 63 |   let (eInp <| outGivenEffect) = f ps 
 64 |   e <- fromCostate handleEffect eInp
 65 |   pure $ outGivenEffect e
 66 |
 67 | ||| Iterates `optimiseStep` `numSteps` times, and logs the progress to the 
 68 | ||| console.  Materialises parameter and state between steps
 69 | public export
 70 | optimise : {p, l : AddCont} -> {e : Cont} ->
 71 |   InterfaceOnPositions l Num =>
 72 |   {default 100 printEvery : Nat} ->
 73 |   Materialise p.Shp => Materialise stateTy =>
 74 |   ScientificDisplay p.Shp => ScientificDisplay l.Shp => ScientificDisplay stateTy =>
 75 |   (f : p =%+> e >-+@ l) ->
 76 |   (handleEffect : Costate (IO <!> e)) ->
 77 |   (initParam : IO p.Shp) ->
 78 |   (opt : Optimiser p stateTy) ->
 79 |   (numSteps : Nat) ->
 80 |   IO (p.Shp, stateTy)
 81 | optimise f handleEffect initParam opt numSteps = do
 82 |   currentValue <- initParam
 83 |   currentState <- opt.initState
 84 |   runActionUntilMaxSteps
 85 |     {l=l.Shp}
 86 |     {printEvery=printEvery}
 87 |     (fromCostate $ optimiseStep f handleEffect opt)
 88 |     numSteps
 89 |     0
 90 |     (currentValue, currentState)
 91 |     (fromCostate $ evalFw (f.fwd . opt.fwd) handleEffect)
 92 |
 93 | ||| TODO is the better name here "buildOptimiser"?
 94 | public export
 95 | buildSupervisedLearningSystem : (f : x =\\=> y) -> (loss : y =\\=> l) ->
 96 |   Materialise (Param f).Shp => InterfaceOnPositions (Param f) Materialise =>
 97 |   Param f =%+> (SupervisedData x.Shp (Param loss).Shp) >-+@ l
 98 | buildSupervisedLearningSystem f loss =
 99 |   let supplied : (a >*< b) >*< c =%+> a >*< (c >*< b)
100 |       supplied = assocL %+>> (id >*< swap)
101 |   in materialiseCont %+>> pushIntoContinuation {d=x>*<Param loss}
102 |        (supplied %+>> (composePara f loss).Run)
103 |
104 |
105 | namespace WithEffect
106 |   ||| Evaluating the total loss over test/inference data in an effectul setting 
107 |   ||| requires a handler for the effect. Usually when the effect is `Dist n`, 
108 |   ||| the handler is simply sampling. We can't do anything else, really!
109 |   public export
110 |   totalLoss : Num l.Shp =>
111 |     (f : x =\\=> e >-+@ y) ->
112 |     (loss : y =\\=> l) ->
113 |     (p : (Param f).Shp) ->
114 |     (handleEffect : Costate (IO <!> e)) ->
115 |     Costate (IO <!> (Const2 (Vect n (x.Shp, (Param loss).Shp)) l.Shp))
116 |   totalLoss (MkPara pCont f) (MkPara z loss) p handleEffect
117 |     = let evalFWithLoss : (x.Shp, z.Shp) -> IO l.Shp
118 |           evalFWithLoss (x, yTrue) = do
119 |             yPred <- fromCostate (evalFw f.fwd handleEffect) (x, p)
120 |             pure $ loss.fwd (yPred, yTrue)
121 |             -- putStrLn "Input: \{show x}, Predicted: \{show yPred}, Loss: \{show lossVal}"
122 |       in toCostate $ \testData => do
123 |         losses <- traverse evalFWithLoss testData
124 |         pure $ Prelude.sum losses
125 |
126 |   ||| Average loss in test/inference 
127 |   public export
128 |   averageLoss :  {n : Nat} ->
129 |     Num l.Shp => Fractional l.Shp => Cast Nat l.Shp =>
130 |     (f : x =\\=> e >-+@ y) ->
131 |     (loss : y =\\=> l) ->
132 |     (p : (Param f).Shp) ->
133 |     (handleEffect : Costate (IO <!> e)) ->
134 |     Costate (IO <!> (Const2 (Vect n (x.Shp, (Param loss).Shp)) l.Shp))
135 |   averageLoss f loss p handleEffect = toCostate $ \testData => do
136 |     lossSum <- fromCostate (totalLoss f loss p handleEffect) testData
137 |     pure (lossSum / cast n)
138 |   
139 | namespace Model
140 |   public export
141 |   train : {a, b : AddCont} -> {l : AddCont} ->
142 |     {default 100 printEvery : Nat} ->
143 |     (m : a -\-> b) ->
144 |     (loss : b =\\=> l) ->
145 |     InterfaceOnPositions l Num =>
146 |     ScientificDisplay l.Shp =>
147 |     Materialise m.Params => Materialise stateTy =>
148 |     ScientificDisplay m.Params => ScientificDisplay stateTy =>
149 |     (trainData : DataLoader a.Shp (Param loss).Shp) ->
150 |     (opt : Optimiser (ParamCont m) stateTy) ->
151 |     (numSteps : Nat) ->
152 |     IO (m.Params, stateTy)
153 |   train m loss trainData opt numSteps = optimise {printEvery}
154 |     (buildSupervisedLearningSystem (toPara m) loss)
155 |     (handleData trainData)
156 |     m.init
157 |     opt
158 |     numSteps
159 |
160 |   ||| Average loss over a dataset
161 |   public export
162 |   averageLoss : {0 a, b, l : AddCont} ->
163 |     (m : a -\-> b) ->
164 |     (loss : b =\\=> l) ->
165 |     Fractional l.Shp => Cast Nat l.Shp =>
166 |     (p : m.Params) ->
167 |     (dl : DataLoader a.Shp (Param loss).Shp) ->
168 |     l.Shp
169 |   averageLoss m (MkPara z loss) p dl =
170 |     let pointLoss : (a.Shp, z.Shp) -> l.Shp
171 |         pointLoss (x, yTrue) = loss.fwd (m.fwd x p, yTrue)
172 |     in Prelude.sum (pointLoss <$> dl.dataset) / cast dl.datasetSize
173 |
174 |   ||| Print a model's predictions on a dataset's inputs
175 |   public export
176 |   evalPrint : {0 a, b : AddCont} ->
177 |     ScientificDisplay a.Shp => ScientificDisplay b.Shp =>
178 |     (m : a -\-> b) -> (p : m.Params) ->
179 |     DataLoader a.Shp b.Shp -> IO ()
180 |   evalPrint m p dl = for_ dl.dataset $ \(x, _) =>
181 |     putStrLn "Input: \{showSci x}, Predicted: \{showSci (m.fwd x p)}"
182 |
183 |
184 | {-
185 | -- todo write a variant of this with effects?
186 | public export
187 | debugPrint : {x, y : AddCont} ->
188 |   Show x.Shp => Show y.Shp =>
189 |   (name : String) ->
190 |   (f : ParaAddMLens {m=IO} x y) ->
191 |   Show (GetParam f).Shp =>
192 |   ParaAddMLens {m=IO} x y
193 | debugPrint name (MkPara pCont f) = MkPara
194 |   pCont
195 |   (!%%+ \(x, p) => do
196 |     putStrLn "--------------------------------"
197 |     putStrLn "\{name} input: \{show x}"
198 |     putStrLn "\{name} parameter: \{show p}"
199 |     (y ** ky) <- (%%!+ f) (x, p)
200 |     putStrLn "\{name} output: \{show y}"
201 |     putStrLn "--------------------------------"
202 |     pure (y ** ky))