src/autograd/sgd.bend

bendtorch-core@0.1.0.0 · 208 lines · 21 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
5import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
6import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
7
8# Plain stochastic gradient descent on the tape, PyTorch's
9# torch.optim.SGD(params, lr) with no momentum, dampening, nesterov,
10# weight_decay or maximize. The parameters are a list of tape ids, each a
11# leaf. step visits them in list order; a parameter whose slot holds g gets
12# the value p - lr * g, cell by cell on the parallel elementwise build
13# (Elementwise.Descend), and keeps its slot, so a later read still sees g; a
14# parameter whose slot is empty, PyTorch's p.grad is None, is left as it was.
15# A listed id twice is updated twice, as PyTorch's step does. Recording an
16# operation after a step reads the new values, as a forward after
17# optimizer.step() reads the updated parameters. Backward.zero_grad empties
18# the slots.
19#
20# step checks every id, then the learning rate, before it updates anything:
21# an id the tape does not hold is refused with UnknownNode, an id of a node
22# that is not a leaf with NotLeaf, and a learning rate below zero (IEEE
23# `lr < 0`, PyTorch's own check: -0.0, NaN and +inf are accepted) with
24# InvalidLearningRate. A refusal leaves the tape unchanged.
25
26type Error is Data:
27  UnknownNode{}
28  NotLeaf{}
29  InvalidLearningRate{}
30  TensorError{error: Core.Error}
31
32def leaf_refusal(node: Backward.Node) -> Maybe<&2, Error>:
33  match node:
34    case Backward.Leaf{}:
35      None{}
36    case Backward.Add{a, b, left, right}:
37      Some{NotLeaf{}}
38    case Backward.Sub{a, b, left, right}:
39      Some{NotLeaf{}}
40    case Backward.Mul{a, b, left, right}:
41      Some{NotLeaf{}}
42    case Backward.Sum{a, input}:
43      Some{NotLeaf{}}
44    case Backward.SumDims{a, input, kept}:
45      Some{NotLeaf{}}
46    case Backward.Matmul{a, b, left, right}:
47      Some{NotLeaf{}}
48    case Backward.MatmulNT{a, b, left, right}:
49      Some{NotLeaf{}}
50    case Backward.Relu{a, input}:
51      Some{NotLeaf{}}
52    case Backward.MseLoss{a, b, left, right}:
53      Some{NotLeaf{}}
54    case Backward.CrossEntropy{a, input}:
55      Some{NotLeaf{}}
56    case Backward.BatchedMatmul{a, b, left, right}:
57      Some{NotLeaf{}}
58    case Backward.Slice{a, input, kept, plan}:
59      Some{NotLeaf{}}
60    case Backward.Gathered{a, input, plan, entries}:
61      Some{NotLeaf{}}
62    case Backward.Constant{}:
63      None{}
64    case Backward.Frozen{id, inner}:
65      leaf_refusal(inner)
66
67# Why id cannot be a parameter of a newest-first list of count nodes, or
68# None when it is a leaf the list holds.
69def refusal_at(nodes: List<&2, Backward.Node>, +count: Nat, +id: Nat, same: Bool) -> Maybe<&2, Error>:
70  match nodes count same:
71    case Nil{} count same:
72      Some{UnknownNode{}}
73    case Con{node, rest} 0n same:
74      Some{UnknownNode{}}
75    case Con{node, rest} 1n+below True{}:
76      leaf_refusal(node)
77    case Con{node, rest} 1n+below False{}:
78      refusal_at(rest, below, id, Backward.at_head(below, id))
79
80def first_refusal(found: Maybe<&2, Error>, later: Maybe<&2, Error>) -> Maybe<&2, Error>:
81  match found:
82    case Some{error}:
83      Some{error}
84    case None{}:
85      later
86
87# The refusal of the first parameter, in list order, that is not a leaf of
88# the tape, or None when every one is.
89def refusal(+nodes: List<&2, Backward.Node>, +count: Nat, params: List<&2, Nat>) -> Maybe<&2, Error>:
90  match params:
91    case Nil{}:
92      None{}
93    case Con{+id, rest}:
94      first_refusal(refusal_at(nodes, count, id, Backward.at_head(count, id)), refusal(nodes, count, rest))
95
96def Walked() -> Type:
97  List<&1, Backward.Entry> & Result<&2, &2, Error, Unit>
98
99def descended(node: Backward.Node, +shape: Shape.Shape, saved: Backward.Saved,
100  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Backward.Entry & Result<&2, &2, Error, Unit>:
101  (operands, result) = pair
102  (value, g) = operands
103  match result:
104    case Fail{error}:
105      (Backward.Entry{node, shape, value, saved, Gradient.Held{g}}, Fail{TensorError{error}})
106    case Done{updated}:
107      (Backward.Entry{node, shape, updated, saved, Gradient.Held{g}}, Done{Unit{}})
108
109# One parameter: p - lr * g, cell by cell, when its slot holds g; unchanged
110# when the slot is empty.
111def descend(+lr: F32, entry: Backward.Entry) -> Backward.Entry & Result<&2, &2, Error, Unit>:
112  Backward.Entry{node, +shape, value, saved, slot} = entry
113  match slot:
114    case Gradient.Unset{}:
115      (Backward.Entry{node, shape, value, saved, Gradient.Unset{}}, Done{Unit{}})
116    case Gradient.Held{g}:
117      descended(node, shape, saved, Elementwise.binary(Elementwise.Descend{lr}, value, g))
118
119def placed(rest: List<&1, Backward.Entry>, pair: Backward.Entry & Result<&2, &2, Error, Unit>) -> Walked():
120  (entry, result) = pair
121  (entry <> rest, result)
122
123def kept(entry: Backward.Entry, pair: Walked()) -> Walked():
124  (rest, result) = pair
125  (entry <> rest, result)
126
127# Updates entry id of a newest-first list of count entries and keeps every
128# other entry.
129def descend_at(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +id: Nat, same: Bool) -> Walked():
130  match entries count same:
131    case Nil{} count same:
132      (Nil{}, Fail{UnknownNode{}})
133    case Con{entry, rest} 0n same:
134      (entry <> rest, Fail{UnknownNode{}})
135    case Con{entry, rest} 1n+below True{}:
136      placed(rest, descend(lr, entry))
137    case Con{entry, rest} 1n+below False{}:
138      kept(entry, descend_at(lr, rest, below, id, Backward.at_head(below, id)))
139
140def continued(rec: List<&1, Backward.Entry> -> Walked(), pair: Walked()) -> Walked():
141  (entries, result) = pair
142  match result:
143    case Fail{error}:
144      (entries, Fail{error})
145    case Done{unit}:
146      rec(entries)
147
148# Each parameter in list order, each on the entries the one before it left.
149def descend_all(+lr: F32, +count: Nat, params: List<&2, Nat>, entries: List<&1, Backward.Entry>) -> Walked():
150  match params:
151    case Nil{}:
152      (entries, Done{Unit{}})
153    case Con{+id, rest}:
154      continued(entries => descend_all(lr, count, rest, entries), descend_at(lr, entries, count, id, Backward.at_head(count, id)))
155
156def walked_tape(pair: Walked()) -> Backward.Tape & Result<&2, &2, Error, Unit>:
157  (entries, result) = pair
158  (Backward.Tape{entries}, result)
159
160def refused_if(below: Bool) -> Maybe<&2, Error>:
161  match below:
162    case True{}:
163      Some{InvalidLearningRate{}}
164    case False{}:
165      None{}
166
167# InvalidLearningRate when lr is below zero, as PyTorch's `lr < 0.0`; None
168# for every other word, -0.0, NaN and +inf included.
169def rate_refusal(+lr: F32) -> Maybe<&2, Error>:
170  refused_if(Elementwise.below_zero(lr))
171
172def rated(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +params: List<&2, Nat>, refused: Maybe<&2, Error>) ->
173  Backward.Tape & Result<&2, &2, Error, Unit>:
174  match refused:
175    case Some{error}:
176      (Backward.Tape{entries}, Fail{error})
177    case None{}:
178      walked_tape(descend_all(lr, count, params, entries))
179
180def checked(+lr: F32, entries: List<&1, Backward.Entry>, +count: Nat, +params: List<&2, Nat>, refused: Maybe<&2, Error>) ->
181  Backward.Tape & Result<&2, &2, Error, Unit>:
182  match refused:
183    case Some{error}:
184      (Backward.Tape{entries}, Fail{error})
185    case None{}:
186      rated(lr, entries, count, params, rate_refusal(lr))
187
188def planned(+lr: F32, +params: List<&2, Nat>, pair: List<&1, Backward.Entry> & List<&2, Backward.Node>) ->
189  Backward.Tape & Result<&2, &2, Error, Unit>:
190  (entries, +nodes) = pair
191  +count = List.length(&2, Backward.Node, nodes)
192  checked(lr, entries, count, params, refusal(nodes, count, params))
193
194# One SGD step with learning rate lr over the parameters params.
195def step(tape: Backward.Tape, +lr: F32, +params: List<&2, Nat>) -> Backward.Tape & Result<&2, &2, Error, Unit>:
196  Backward.Tape{entries} = tape
197  planned(lr, params, Backward.listed(entries))
198
199def error_code(error: Error) -> U32:
200  match error:
201    case UnknownNode{}:
202      Backward.error_code(Backward.UnknownNode{})
203    case NotLeaf{}:
204      109
205    case InvalidLearningRate{}:
206      110
207    case TensorError{error}:
208      Core.error_code(error)