0 | module NN.Training.DataLoader
 1 |
 2 | import Data.Tensor
 3 | import Data.Container.Additive
 4 |
 5 | import Control.Monad.Distribution
 6 | import Control.Monad.Sample.Definition
 7 | import Control.Monad.Sample.Instances
 8 |
 9 | public export
10 | DatasetAxis : Axis
11 | DatasetAxis = "dataset" ~> List1
12 |
13 | ||| A container whose shapes are tensors over a List1 container, and 
14 | ||| positions are specific choices of input-output pairs
15 | ||| TODO Does any other container other than `List1` make sense?
16 | public export
17 | DataLoader : (input, output : Type) -> Cont
18 | DataLoader input output = Pick [DatasetAxis] (input, output)
19 |
20 | public export
21 | Dataset : (input, output : Type) -> Type
22 | Dataset input output = (DataLoader input output).Shp
23 |
24 | ||| The number of entries, one more than the shape of `List1`
25 | public export
26 | datasetSize : Dataset input output -> Nat
27 | datasetSize dl = S dl.extractShapeRank1
28 |
29 | ||| On the forward pass: given a dataset, get its size, and produce a uniform 
30 | ||| distribution over that many elements.
31 | ||| On the bw pass: given the index in that size, pick the actual element 
32 | ||| corresponding to it
33 | public export
34 | dataLens : DataLoader input output =%> Dist DatasetAxis
35 | dataLens = Tensor.pick %>> rightUnit %>> uniformLens
36 |
37 | ||| Handler for the dataset, samples from it uniformly
38 | public export
39 | handleData : Costate (IO <!> DataLoader input output)
40 | handleData = (IO <!> dataLens) %>> sample
41 |
42 | ||| The container used to store data for a supervised learning system
43 | ||| Single shape, positions are (input, output) pairs
44 | public export
45 | SupervisedData : (input, output : Type) -> Cont
46 | SupervisedData input output = Nap (input, output)
47 |
48 | ||| Sampling from one particular dataset
49 | public export
50 | handleDataFor : (dl : Dataset x y) -> Costate (IO <!> SupervisedData x y)
51 | handleDataFor dl = (IO <!> atShape {c = DataLoader x y} dl) %>> handleData
52 |
53 | ||| A dataset from a batch of inputs and a ground-truth function
54 | public export
55 | makeDataLoader : Monad m =>
56 |   {batch : Axis} -> (isCub : IsCubical batch) => IsSucc (batch.dim) =>
57 |   (xs : Tensor [batch] input) ->
58 |   (groundTruthFn : input -> m output) ->
59 |   m (Dataset input output)
60 | makeDataLoader {isCub = MkIsCubical _ (S k)} xs groundTruthFn
61 |   = restructureAxis vectToList1 . liftA2Tensor xs <$> traverse groundTruthFn xs