0 | module NN.Training.Training
  1 |
  2 | import Data.Either
  3 |
  4 | import Data.Tensor
  5 | import Data.Container.Additive as Additive
  6 | import public Data.ScientificNotation
  7 | import NN.Optimisers
  8 |
  9 | import NN.Utils
 10 | import NN.Training.DataLoader
 11 | import NN.Architectures.LossFunctions
 12 | import Data.Para
 13 | import Data.Autodiff.Model
 14 |
 15 | {-------------------------------------------------------------------------------
 16 | {-------------------------------------------------------------------------------
 17 | TODO update this in light of effects
 18 |
 19 | This file defines functions which perform pure optimisation:
 20 | optimisation of a differentiable function `f : p -> x` (here modelled as a lens `f : p =%> x` via some optimiser such as gradient descent. 
 21 | Here no "loss function" or "input-output pairs" are needed, just a function to optimise.
 22 |
 23 | This file provides functionality for creating, and turning supervised learning problems into pure optimisation problems, via a function which takes:
 24 | a) a parametric lens `f : x >< p =%>> y`
 25 | b) a loss function `loss : (y, y) =%> l`
 26 | c) input-output pairs `IO (x.Shp, y.Shp)`
 27 | and composes them to produce an optimisation problem `f : p =%> l` which the above described functions can consume.
 28 |
 29 | 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.
 30 |
 31 |
 32 | todo using Hom-version of optimisation becomes problematic if we either
 33 | a) have the dependency of the parameter on the input
 34 | b) use monadic lenses
 35 |
 36 | -------------------------------------------------------------------------------}
 37 | -------------------------------------------------------------------------------}
 38 |
 39 | ||| Performs a single step of optimisation of some differentiable function
 40 | ||| `f : p -> l`, additionally handling some effect `e`
 41 | ||| The optimiser used is allowed to be stateful meaning the result of the
 42 | ||| optimisation is both the final parameter and the state of the optimiser
 43 | public export
 44 | optimiseStep : {p, l : AddCont} -> {e : Cont} ->
 45 |   InterfaceOnPositions l Num =>
 46 |   (f : p =%+> e >-+@ l) ->
 47 |   (handleEffect : Costate (IO <!> e)) ->
 48 |   (optimiser : Optimiser p stateTy) ->
 49 |   Costate (IO <!> (Const (p.Shp, stateTy)))
 50 | optimiseStep f handleEffect (MkOptimiser opt _) =
 51 |   let closeFunction : p =%+> e >-+@ Scalar
 52 |       closeFunction = f %+>> (id >-+@ constantOne)
 53 |
 54 |       closeFunctionT : UC p =%> e
 55 |       closeFunctionT = adjL closeFunction
 56 |   in (IO <!> (opt %>> closeFunctionT)) %>> handleEffect
 57 |
 58 | ||| Evaluates the forward pass of some effectful lens
 59 | public export
 60 | evalFw : {0 e : Cont} ->
 61 |   (f : a -> Ext e b) ->
 62 |   (handleEffect : Costate (IO <!> e)) ->
 63 |   Costate (IO <!> (Const2 a b))
 64 | evalFw f handleEffect = toCostate $ \ps => do
 65 |   let (eInp <| outGivenEffect) = f ps 
 66 |   e <- fromCostate handleEffect eInp
 67 |   pure $ outGivenEffect e
 68 |
 69 | ||| Iterates `optimiseStep` `numSteps` times, and logs the progress to the 
 70 | ||| console.  Materialises parameter and state between steps
 71 | public export
 72 | optimise : {p, l : AddCont} -> {e : Cont} ->
 73 |   InterfaceOnPositions l Num =>
 74 |   {default 100 printEvery : Nat} ->
 75 |   Materialise p.Shp => Materialise stateTy =>
 76 |   ScientificDisplay p.Shp => ScientificDisplay l.Shp => ScientificDisplay stateTy =>
 77 |   (f : p =%+> e >-+@ l) ->
 78 |   (handleEffect : Costate (IO <!> e)) ->
 79 |   (initParam : IO p.Shp) ->
 80 |   (opt : Optimiser p stateTy) ->
 81 |   (numSteps : Nat) ->
 82 |   IO (p.Shp, stateTy)
 83 | optimise f handleEffect initParam opt numSteps = do
 84 |   currentValue <- initParam
 85 |   currentState <- opt.initState
 86 |   runActionUntilMaxSteps
 87 |     {l=l.Shp}
 88 |     {printEvery=printEvery}
 89 |     (fromCostate $ optimiseStep f handleEffect opt)
 90 |     numSteps
 91 |     0
 92 |     (currentValue, currentState)
 93 |     (fromCostate $ evalFw (f.fwd . opt.fwd) handleEffect)
 94 |
 95 | ||| TODO is the better name here "buildOptimiser"?
 96 | ||| TODO should now in the presence of `DataLoader` container `SupervisedData` 
 97 | ||| be replaced with it?
 98 | public export
 99 | buildSupervisedLearningSystem : (f : x =\\=> y) -> (loss : y =\\=> l) ->
100 |   Materialise (Param f).Shp => InterfaceOnPositions (Param f) Materialise =>
101 |   Param f =%+> (SupervisedData x.Shp (Param loss).Shp) >-+@ l
102 | buildSupervisedLearningSystem f loss =
103 |   let supplied : (a >*< b) >*< c =%+> a >*< (c >*< b)
104 |       supplied = assocL %+>> (id >*< swap)
105 |   in materialiseCont %+>> pushIntoContinuation {d=x>*<Param loss}
106 |        (supplied %+>> (composePara f loss).Run)
107 |
108 |
109 | namespace WithEffect
110 |   ||| Computes loss for one input-output pair
111 |   public export
112 |   computeLoss :
113 |     (f : x =\\=> e >-+@ y) ->
114 |     (loss : y =\\=> l) ->
115 |     (p : (Param f).Shp) ->
116 |     (handleEffect : Costate (IO <!> e)) ->
117 |     Costate (IO <!> Const2 (x.Shp, (Param loss).Shp) l.Shp)
118 |   computeLoss f loss p handleEffect = toCostate $ \(x, yTrue) => do
119 |     yPred <- fromCostate (evalFw (Run f).fwd handleEffect) (x, p)
120 |     pure ((Run loss).fwd (yPred, yTrue))
121 |     
122 |
123 |
124 |   ||| Evaluating the total loss over test/inference data in an effectul setting 
125 |   ||| requires a handler for the effect. Usually when the effect is `Dist n`, 
126 |   ||| the handler is simply sampling. We can't do anything else, really!
127 |   public export
128 |   totalLoss : Num l.Shp =>
129 |     (f : x =\\=> e >-+@ y) ->
130 |     (loss : y =\\=> l) ->
131 |     (p : (Param f).Shp) ->
132 |     (handleEffect : Costate (IO <!> e)) ->
133 |     Costate (IO <!> (Const2 (Vect n (x.Shp, Label loss)) l.Shp))
134 |   totalLoss (MkPara pCont f) (MkPara z loss) p handleEffect
135 |     = let evalFWithLoss : (x.Shp, z.Shp) -> IO l.Shp
136 |           evalFWithLoss (x, yTrue) = do
137 |             yPred <- fromCostate (evalFw f.fwd handleEffect) (x, p)
138 |             pure $ loss.fwd (yPred, yTrue)
139 |             -- putStrLn "Input: \{show x}, Predicted: \{show yPred}, Loss: \{show lossVal}"
140 |       in toCostate $ \testData => do
141 |         losses <- traverse evalFWithLoss testData
142 |         pure $ Prelude.sum losses
143 |
144 |   ||| Average loss in test/inference 
145 |   ||| What happens when branches are mismatched?
146 |   public export
147 |   averageLoss :  {n : Nat} ->
148 |     Num l.Shp => Fractional l.Shp => Cast Nat l.Shp =>
149 |     (f : x =\\=> e >-+@ y) ->
150 |     (loss : y =\\=> l) ->
151 |     (p : (Param f).Shp) ->
152 |     (handleEffect : Costate (IO <!> e)) ->
153 |     Costate (IO <!> (Const2 (Vect n (x.Shp, Label loss)) l.Shp))
154 |   averageLoss f loss p handleEffect = toCostate $ \testData => do
155 |     lossSum <- fromCostate (totalLoss f loss p handleEffect) testData
156 |     pure (lossSum / cast n)
157 |
158 |   ||| Print a model's predictions on a dataset's inputs, handling its effect.
159 |   ||| By default it also prints the shape of the effect
160 |   public export
161 |   evalPrint : {0 a, y : AddCont} -> {0 e : Cont} ->
162 |     {default True printEffect : Bool} ->
163 |     Show a.Shp => Show e.Shp => Show y.Shp =>
164 |     (m : a -\-> e >-+@ y) -> (p : m.Params) ->
165 |     (handleEffect : Costate (IO <!> e)) ->
166 |     Dataset a.Shp lab -> IO ()
167 |   evalPrint {printEffect} m p handleEffect dl = for_ dl $ \(x, _) => do
168 |     let out = m.fwd x p
169 |     answer <- fromCostate handleEffect (shapeExt out)
170 |     let asked : String -- without the type annotation there's an error
171 |         asked = if printEffect then ", Effect out: \{show (shapeExt out)}" else ""
172 |     putStrLn "Input: \{show x}\{asked}, Predicted: \{show (index out answer)}"
173 |
174 |   public export
175 |   evalLoss : {0 a, y, l : AddCont} -> {0 e : Cont} ->
176 |     (m : a -\-> e >-+@ y) -> (loss : Loss y l) -> (p : m.Params) ->
177 |     (handleEffect : Costate (IO <!> e)) ->
178 |     Dataset a.Shp (Label loss) -> IO (Tensor [DatasetAxis] l.Shp)
179 |   evalLoss m loss p handleEffect
180 |     = traverse (fromCostate (computeLoss (toPara m) loss p handleEffect))
181 |
182 |   ||| The fraction of inputs a partial loss was defined on
183 |   public export
184 |   matchRate : Foldable t => t (Either e l) -> Double
185 |   matchRate xs = cast (count isRight xs) / cast (count (const True) xs)
186 |
187 |   ||| The mean of a partial loss over the entries it was defined on
188 |   public export
189 |   averageMatched : Foldable t => Fractional l => Cast Nat l =>
190 |     t (Either e l) -> Maybe l
191 |   averageMatched xs = case rights (toList xs) of
192 |     [] => Nothing
193 |     ls => Just (Prelude.sum ls / cast (length ls))
194 |   
195 | namespace Model
196 |   public export
197 |   train : {a, b : AddCont} -> {l : AddCont} ->
198 |     {default 100 printEvery : Nat} ->
199 |     (m : a -\-> b) ->
200 |     (loss : Loss b l) ->
201 |     InterfaceOnPositions l Num =>
202 |     ScientificDisplay l.Shp =>
203 |     Materialise m.Params => Materialise stateTy =>
204 |     ScientificDisplay m.Params => ScientificDisplay stateTy =>
205 |     (trainData : Dataset a.Shp (Label loss)) ->
206 |     (opt : Optimiser (ParamCont m) stateTy) ->
207 |     (numSteps : Nat) ->
208 |     IO (m.Params, stateTy)
209 |   train m loss trainData opt numSteps = optimise {printEvery}
210 |     (buildSupervisedLearningSystem (toPara m) loss)
211 |     (handleDataFor trainData)
212 |     m.init
213 |     opt
214 |     numSteps
215 |
216 |   ||| Average loss over a dataset
217 |   public export
218 |   averageLoss : {0 a, b, l : AddCont} ->
219 |     (m : a -\-> b) ->
220 |     (loss : b =\\=> l) ->
221 |     Fractional l.Shp => Cast Nat l.Shp =>
222 |     (p : m.Params) ->
223 |     (dl : Dataset a.Shp (Label loss)) ->
224 |     l.Shp
225 |   averageLoss m (MkPara z loss) p dl =
226 |     let pointLoss : (a.Shp, z.Shp) -> l.Shp
227 |         pointLoss (x, yTrue) = loss.fwd (m.fwd x p, yTrue)
228 |     in Prelude.sum (pointLoss <$> dl) / cast (datasetSize dl)
229 |
230 |   ||| Print model's predictions on dataset's inputs
231 |   ||| TODO make this be printed in a more pretty way
232 |   public export
233 |   evalPrint : {0 a, b : AddCont} -> {dataset : Axis} ->
234 |     IsFoldable dataset.cont =>
235 |     ScientificDisplay a.Shp => ScientificDisplay b.Shp =>
236 |     (m : a -\-> b) -> (p : m.Params) ->
237 |     Tensor [dataset] a.Shp -> IO ()
238 |   evalPrint m p dl = do
239 |     putStrLn "Model's predictions on inputs:"
240 |     let input = "Input"
241 |         predicted = "Predicted"
242 |     for_ dl $ \x =>
243 |       putStrLn "  \{dim input}: \{showSci x}, \{dim predicted}: \{bold $ showSci (m.fwd x p)}"
244 |
245 | {-
246 | -- todo write a variant of this with effects?
247 | public export
248 | debugPrint : {x, y : AddCont} ->
249 |   Show x.Shp => Show y.Shp =>
250 |   (name : String) ->
251 |   (f : ParaAddMLens {m=IO} x y) ->
252 |   Show (GetParam f).Shp =>
253 |   ParaAddMLens {m=IO} x y
254 | debugPrint name (MkPara pCont f) = MkPara
255 |   pCont
256 |   (!%%+ \(x, p) => do
257 |     putStrLn "--------------------------------"
258 |     putStrLn "\{name} input: \{show x}"
259 |     putStrLn "\{name} parameter: \{show p}"
260 |     (y ** ky) <- (%%!+ f) (x, p)
261 |     putStrLn "\{name} output: \{show y}"
262 |     putStrLn "--------------------------------"
263 |     pure (y ** ky))