tensor_model.bend

stelliferous@0.0.2.0 · 170 lines · 32 definitions

1# Models: a step computes with tensors it takes by name from a store, so a
2# network reads as one step after another and its weights are loaded once.
3import Base
4import stelliferous@0.0.2.0/tensor.bend as Tensors
5import stelliferous@0.0.2.0/tensor_layout.bend as Layout
6import stelliferous@0.0.2.0/tensor_file.bend as Files
7
8def Tensor() -> Type:
9  Tensors.Tensor<F32>
10
11type Store is Type:
12  Empty{}
13  Entry{name: String,tensor: Tensor(),rest: Store}
14
15# A step: given the store, a value and the store that remains.
16def Model(A: Type) -> Type:
17  Store -> A & Store
18
19def Model.pure(-A: Type,value: A) -> Model(A):
20  store => (value,store)
21
22def Model.next(-A: Type,-B: Type,pair: A & Store,next: A -> Model(B)) -> B & Store:
23  (value,store) = pair
24  next(value)(store)
25
26def Model.bind(-A: Type,-B: Type,step: Model(A),next: A -> Model(B)) -> Model(B):
27  store => Model.next(A,B,step(store),next)
28
29def missing(+name: String) -> Tensor():
30  Tensors.failed(~F32,"take",Tensors.MissingName{name},[],0.0,1)
31
32def passed(name: String,tensor: Tensor(),pair: Tensor() & Store) -> Tensor() & Store:
33  (result,store) = pair
34  (result,Entry{name,tensor,store})
35
36def scan(rest: Store,same: Bool,+name: String,tensor: Tensor(),+wanted: String) -> Tensor() & Store:
37  match rest same:
38    case Empty{} True{}: (tensor,Empty{})
39    case Entry{next,other,more} True{}: (tensor,Entry{next,other,more})
40    case Empty{} False{}: (missing(wanted),Entry{name,tensor,Empty{}})
41    case Entry{+next,other,more} False{}: passed(name,tensor,scan(more,String.eq(next,wanted),next,other,wanted))
42
43def taken(store: Store,+wanted: String) -> Tensor() & Store:
44  match store:
45    case Empty{}: (missing(wanted),Empty{})
46    case Entry{+name,tensor,rest}: scan(rest,String.eq(name,wanted),name,tensor,wanted)
47
48# The tensor of that name leaves the store. A name that is not there gives a
49# failed tensor.
50def take(+name: String) -> Model(Tensor()):
51  store => taken(store,name)
52
53def kept(+name: String,pair: Tensor() & Tensor(),store: Store) -> Tensor() & Store:
54  (tensor,copy) = pair
55  (tensor,Entry{name,copy,store})
56
57# Puts a copy under the name, for a later take, and gives the tensor back.
58def keep(+name: String,tensor: Tensor()) -> Model(Tensor()):
59  store => kept(name,Layout.copy(~F32,tensor),store)
60
61def run(-A: Type,step: Model(A),store: Store) -> A & Store:
62  step(store)
63
64def value(-A: Type,pair: A & Store) -> A:
65  (result,store) = pair
66  result
67
68# Runs the steps on the store and gives the result; what remains in the store is dropped.
69def result(-A: Type,step: Model(A),store: Store) -> A:
70  value(A,step(store))
71
72def joined(+name: String,first: Tensor(),second: Tensor(),rest: Store & Store) -> Store & Store:
73  (left,right) = rest
74  (Entry{name,first,left},Entry{name,second,right})
75
76def copied_entry(+name: String,pair: Tensor() & Tensor(),rest: Store & Store) -> Store & Store:
77  (first,second) = pair
78  joined(name,first,second,rest)
79
80# Two stores with the same tensors, to run a model more than once.
81def copy(store: Store) -> Store & Store:
82  match store:
83    case Empty{}: (Empty{},Empty{})
84    case Entry{+name,tensor,rest}: copied_entry(name,Layout.copy(~F32,tensor),copy(rest))
85
86# Loading: folder/names.txt lists one name per line; folder/<name>.npy holds its tensor.
87
88def entry(+name: String,tensor: Tensor(),store: Store) -> Store:
89  Entry{name,tensor,store}
90
91def load_each(names: List<&2,String>,+folder: String) -> IO(Store):
92  match names:
93    case Nil{}: IO.pure(Store,Empty{})
94    case Con{+name,rest}:
95      match name:
96        case SNil{}: load_each(rest,folder)
97        case SCon{head,tail}:
98          do IO<Store>:
99            tensor : Tensor() <- Files.load_npy(folder ++ "/" ++ name ++ ".npy")
100            store : Store <- load_each(rest,folder)
101            IO.pure(Store,entry(name,tensor,store))
102
103def named(text: Result<&1,&1,U32 & String,String>,+folder: String) -> IO(Store):
104  match text:
105    case Fail{error}: IO.pure(Store,Empty{})
106    case Done{text}: load_each(String.lines(text),folder)
107
108def listed(read: File & Result<&1,&1,U32 & String,String>,+folder: String) -> IO(Store):
109  (file,text) = read
110  do IO<Store>:
111    File.close(file)
112    named(text,folder)
113
114def opened(result: Result<&1,&1,U32 & String,File>,+folder: String) -> IO(Store):
115  match result:
116    case Fail{error}: IO.pure(Store,Empty{})
117    case Done{file}:
118      do IO<Store>:
119        read : File & Result<&1,&1,U32 & String,String> <- File.read(file,1048576)
120        listed(read,folder)
121
122# A folder without names.txt gives an empty store; every take then fails by name.
123def load(+folder: String) -> IO(Store):
124  do IO<Store>:
125    result : Result<&1,&1,U32 & String,File> <- File.open(folder ++ "/names.txt","r")
126    opened(result,folder)
127
128# Saving: the tensors of a store as folder/<name>.npy, and their names in folder/names.txt.
129
130def named_rest(+name: String,tensor: Tensor(),pair: Store & String) -> Store & String:
131  (rest,text) = pair
132  (Entry{name,tensor,rest},name ++ "
133" ++ text)
134
135def first(written: File & Result<&1,&1,U32 & String,Unit>) -> File:
136  (file,result) = written
137  file
138
139def names(store: Store) -> Store & String:
140  match store:
141    case Empty{}: (Empty{},"")
142    case Entry{+name,tensor,rest}: named_rest(name,tensor,names(rest))
143
144def save_each(store: Store,+folder: String) -> IO(Store):
145  match store:
146    case Empty{}: IO.pure(Store,Empty{})
147    case Entry{+name,tensor,rest}:
148      do IO<Store>:
149        saved : Tensor() <- Files.save_npy(folder ++ "/" ++ name ++ ".npy",tensor)
150        others : Store <- save_each(rest,folder)
151        IO.pure(Store,entry(name,saved,others))
152
153def listing(result: Result<&1,&1,U32 & String,File>,text: String) -> IO(Unit):
154  match result:
155    case Fail{error}: IO.pure(Unit,Unit{})
156    case Done{file}:
157      do IO<Unit>:
158        written : File & Result<&1,&1,U32 & String,Unit> <- File.write(file,text)
159        File.close(first(written))
160
161def save_listed(pair: Store & String,+folder: String) -> IO(Store):
162  (store,text) = pair
163  do IO<Store>:
164    result : Result<&1,&1,U32 & String,File> <- File.open(folder ++ "/names.txt","w")
165    listing(result,text)
166    save_each(store,folder)
167
168# Gives the store back; a tensor that could not be written comes back failed.
169def save(+folder: String,store: Store) -> IO(Store):
170  save_listed(names(store),folder)