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 | import NN.Architectures.LossFunctions
  7 |
  8 | import NN.Utils
  9 | import NN.Training.DataLoader
 10 | import Data.Para
 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} -> InterfaceOnPositions l Num =>
 42 |   (f : p =%+> e >-+@ l) ->
 43 |   (handleEffect : Costate (IO <!> e)) ->
 44 |   (optimiser : Optimiser p stateTy) ->
 45 |   Costate (IO <!> (Const (p.Shp, stateTy)))
 46 | optimiseStep f handleEffect (MkOptimiser opt _ _) = 
 47 |   let closeFunction : p =%+> !* e
 48 |       closeFunction = f %+>> (id >-+@ constantOne) %+>> actionToFree
 49 |
 50 |       closeFunctionT : UC p =%> e -- transposed variant of closeFunction
 51 |       closeFunctionT = addContTransposeInv closeFunction
 52 |
 53 |   in (IO <!> (opt %>> closeFunctionT)) %>> handleEffect
 54 |
 55 | ||| Evaluates the forward pass of some effectful lens
 56 | public export
 57 | evalFw : {0 e : Cont} ->
 58 |   (f : a -> Ext e b) ->
 59 |   (handleEffect : Costate (IO <!> e)) ->
 60 |   Costate (IO <!> (Const2 a b))
 61 | evalFw f handleEffect = toCostate $ \ps => do
 62 |   let (eInp <| outGivenEffect) = f ps 
 63 |   e <- fromCostate handleEffect eInp
 64 |   pure $ outGivenEffect e
 65 |
 66 | ||| Iterates `optimiseStep` `numSteps` times, and logs the progress to the 
 67 | ||| console
 68 | public export
 69 | optimise : {p, l : AddCont} -> {e : Cont} -> InterfaceOnPositions l Num =>
 70 |   {default 100 printEvery : Nat} ->
 71 |   {default Nothing customInitParam : Maybe p.Shp} ->
 72 |   ScientificDisplay p.Shp => ScientificDisplay l.Shp => ScientificDisplay stateTy =>
 73 |   (f : p =%+> e >-+@ l) ->
 74 |   (handleEffect : Costate (IO <!> e)) ->
 75 |   (opt : Optimiser p stateTy) ->
 76 |   (numSteps : Nat) ->
 77 |   IO (p.Shp, stateTy)
 78 | optimise f handleEffect opt numSteps = do
 79 |   currentValue : p.Shp <- case customInitParam of
 80 |     Just p => pure p
 81 |     Nothing => opt.initParam
 82 |   currentState <- opt.initState
 83 |   runActionUntilMaxSteps
 84 |     {l=l.Shp}
 85 |     {printEvery=printEvery}
 86 |     (fromCostate $ optimiseStep f handleEffect opt)
 87 |     numSteps
 88 |     0
 89 |     (currentValue, currentState)
 90 |     (fromCostate $ evalFw (f.fwd . opt.fwd) handleEffect)
 91 |
 92 | ||| Given
 93 | ||| a) a parametric lens `f : x >< p =%+> y`
 94 | ||| b) a loss function `loss : y >< y =%+> l`
 95 | ||| builds an effectful lens `p =%+> l`
 96 | public export
 97 | buildSupervisedLearningSystem : (f : ParaAddDLens x y) ->
 98 |   (loss : Loss y {l=l}) ->
 99 |   (GetParam f) =%+> SupervisedData x.Shp y.Shp >-+@ l
100 | buildSupervisedLearningSystem (MkPara p f) loss =
101 |   let rebracket : ((x >*< y) >*< p) =%+> ((x >*< p) >*< y)
102 |       rebracket = assocL %+>> (id >*< swap) %+>> assocR
103 |   in pushIntoContinuation {d=x>*<y} (rebracket %+>> (f >*< id) %+>> loss)
104 |
105 |
106 | namespace WithEffect
107 |   ||| When it comes to effects which involve sampling, where the 'correct' answer
108 |   ||| is stored in the test data, there are different ways of evaluating the loss
109 |   ||| One is to use that correct label to force the correct branch to run, but
110 |   ||| that is impossible with the current type signature
111 |   ||| Instead, we opt out for the more accurate method of sampling during loss
112 |   ||| evaluation
113 |   public export
114 |   totalLoss : Show l.Shp => Num l.Shp =>
115 |     (f : ParaAddDLens x (e >-+@ y)) ->
116 |     (loss : Loss y {l=l}) ->
117 |     (p : (GetParam f).Shp) ->
118 |     (handleEffect : Costate (IO <!> e)) ->
119 |     Costate (IO <!> (Const2 (Vect n (x.Shp, y.Shp)) l.Shp))
120 |   totalLoss (MkPara pCont f) loss p handleEffect = toCostate $ \testData => do
121 |     let evalFWithLoss : (x.Shp, y.Shp) -> IO l.Shp
122 |         evalFWithLoss (x, yTrue) = do
123 |           yPred <- fromCostate (evalFw f.fwd handleEffect) (x, p)
124 |           pure $ loss.fwd (yPred, yTrue)
125 |           -- putStrLn "Input: \{show x}, Predicted: \{show yPred}, Loss: \{show lossVal}"
126 |     losses <- traverse evalFWithLoss testData
127 |     pure $ Prelude.sum losses
128 |   
129 |   public export
130 |   averageLoss :  {n : Nat} ->
131 |     Show l.Shp => Num l.Shp => Fractional l.Shp => Cast Nat l.Shp =>
132 |     (f : ParaAddDLens x (e >-+@ y)) ->
133 |     (loss : Loss y {l=l}) ->
134 |     (p : (GetParam f).Shp) ->
135 |     (handleEffect : Costate (IO <!> e)) ->
136 |     Costate (IO <!> (Const2 (Vect n (x.Shp, y.Shp)) l.Shp))
137 |   averageLoss f loss p handleEffect = toCostate $ \testData => do
138 |     lossSum <- fromCostate (totalLoss f loss p handleEffect) testData
139 |     pure (lossSum / cast n)
140 |   
141 |   ||| Eval a model and loss with specific parameters, in the presence of an effect
142 |   public export
143 |   evalWithLoss : ScientificDisplay x.Shp => ScientificDisplay y.Shp => ScientificDisplay l.Shp =>
144 |     (f : ParaAddDLens x (e >-+@ y)) ->
145 |     (loss : Loss y {l=l}) ->
146 |     (p : (GetParam f).Shp) ->
147 |     (handleEffect : Costate (IO <!> e)) ->
148 |     Costate (IO <!> (Const2 (Vect n (x.Shp, y.Shp)) Unit))
149 |   evalWithLoss (MkPara pCont f) loss p handleEffect = toCostate $ \testData => do 
150 |     let evalFWithLoss : (x.Shp, y.Shp) -> IO ()
151 |         evalFWithLoss (x, yTrue) = do
152 |           yPred <- fromCostate (evalFw f.fwd handleEffect) (x, p)
153 |           let lossVal = loss.fwd (yPred, yTrue)
154 |           putStrLn "Input: \{showSci x}, Predicted: \{showSci yPred}, Loss: \{showSci lossVal}"
155 |     _ <- traverse evalFWithLoss testData
156 |     pure ()
157 |   
158 |   ||| Eval a model with specific parameters, in the presence of an effect
159 |   public export
160 |   eval : ScientificDisplay x.Shp => ScientificDisplay y.Shp =>
161 |     (f : ParaAddDLens x (e >-+@ y)) ->
162 |     (p : (GetParam f).Shp) ->
163 |     (handleEffect : Costate (IO <!> e)) ->
164 |     Costate (IO <!> (Const2 (Vect n x.Shp) Unit))
165 |   eval (MkPara _ f) p handleEffect = toCostate $ \testData => do
166 |     let evalF : x.Shp -> IO ()
167 |         evalF x = do
168 |           yPred <- fromCostate (evalFw f.fwd handleEffect) (x, p)
169 |           putStrLn "Input: \{showSci x}, Predicted: \{showSci yPred}"
170 |     _ <- traverse evalF testData
171 |     pure ()
172 |
173 | namespace WithoutEffect
174 |   public export
175 |   trivialEffect : {y : AddCont} ->
176 |     ParaAddDLens x y -> ParaAddDLens x (Scalar >+@ y)
177 |   trivialEffect (MkPara p f) = MkPara p
178 |     (f %+>> leftUnitInv)
179 |
180 |   public export
181 |   handleTrivial : Costate (IO <!> Scalar)
182 |   handleTrivial = toCostate $ \() => pure ()
183 |
184 |
185 |   ||| Eval a model with specific parameters
186 |   public export
187 |   eval : {y : AddCont} -> ScientificDisplay x.Shp => ScientificDisplay y.Shp =>
188 |     (f : ParaAddDLens x y) ->
189 |     (p : (GetParam f).Shp) ->
190 |     Costate (IO <!> (Const2 (Vect n x.Shp) Unit))
191 |   eval (MkPara pCont f) p
192 |     = eval {e=Scalar} (MkPara pCont (f %+>> unitor)) p handleTrivial
193 |
194 |   public export
195 |   averageLoss :  {y : AddCont} -> {n : Nat} ->
196 |     Show l.Shp => Num l.Shp => Fractional l.Shp => Cast Nat l.Shp =>
197 |     (f : ParaAddDLens x y) ->
198 |     (loss : Loss y {l=l}) ->
199 |     (p : (GetParam f).Shp) ->
200 |     Costate (IO <!> (Const2 (Vect n (x.Shp, y.Shp)) l.Shp))
201 |   averageLoss (MkPara pCont f) loss p = averageLoss {e=Scalar}
202 |     (MkPara pCont (f %+>> unitor))
203 |     loss
204 |     p
205 |     handleTrivial
206 |   
207 |
208 | {-
209 | public export
210 | train : {x, y, l : AddCont} -> InterfaceOnPositions l Num => IsFlat l =>
211 |   {default 100 printEvery : Nat} ->
212 |   (f : ParaAddDLens x y) ->
213 |   Show (GetParam f).Shp => Num l.Shp =>
214 |   Show x.Shp => Show y.Shp => Show stateTy => Show l.Shp =>
215 |   {default Nothing initParam : Maybe (GetParam f).Shp} ->
216 |   (loss : (y >< y) =%> l) ->
217 |   (handleData : Costate (IO <!> (pushDown (x >< y)))) ->
218 |   (opt : Optimiser (GetParam f) stateTy) ->
219 |   (numSteps : Nat) ->
220 |   IO ((GetParam f).Shp, stateTy)
221 | train f loss handleData = optimise
222 |   {e=pushDown (x><y)}
223 |   {printEvery=printEvery}
224 |   {initParam=initParam}
225 |   (buildSupervisedLearningSystem f loss)
226 |   handleData
227 | -}
228 |
229 | {-
230 | -- todo write a variant of this with effects?
231 | public export
232 | debugPrint : {x, y : AddCont} ->
233 |   Show x.Shp => Show y.Shp =>
234 |   (name : String) ->
235 |   (f : ParaAddMLens {m=IO} x y) ->
236 |   Show (GetParam f).Shp =>
237 |   ParaAddMLens {m=IO} x y
238 | debugPrint name (MkPara pCont f) = MkPara
239 |   pCont
240 |   (!%%+ \(x, p) => do
241 |     putStrLn "--------------------------------"
242 |     putStrLn "\{name} input: \{show x}"
243 |     putStrLn "\{name} parameter: \{show p}"
244 |     (y ** ky) <- (%%!+ f) (x, p)
245 |     putStrLn "\{name} output: \{show y}"
246 |     putStrLn "--------------------------------"
247 |     pure (y ** ky))
248 |
249 | -- namespace Additive
250 | --   ||| Evaluates a the forward pass of some effectful lens
251 | --   public export
252 | --   evalFw : {0 a, e, b : AddCont} ->
253 | --     (f : a =%> (e >@ b)) ->
254 | --     (handleEffect : Costate (IO <!> e)) ->
255 | --     Costate (IO <!> (Const2 a.Shp b.Shp))
256 | --   evalFw f handleEffect = toCostate $ \ps => do
257 | --     let (eInp <| outGivenEffect) = f.fwd ps 
258 | --     e <- fromCostate handleEffect eInp
259 | --     pure $ outGivenEffect e
260 |