tensor_training.bend
stelliferous@0.0.3.0 · 218 lines · 0 definitions
1# Stelliferous training: convolution, batch norm, activation and pool layers.
2# A layer is a value. forward returns the output, what backward reads and the
3# layer, backward the input gradient and the layer holding its parameter
4# gradients, evaluate the output in inference mode, and sgd_step the layer with
5# its gradients applied. Every step is a public operation in a fixed FP32
6# order, so results do not depend on the worker count.
7import Base
8import stelliferous@0.0.3.0/math_activation.bend as Activations
9import stelliferous@0.0.3.0/tensor.bend as Tensors
10import stelliferous@0.0.3.0/tensor_f32.bend as F
11import stelliferous@0.0.3.0/tensor_layer.bend as ConvLayer
12import stelliferous@0.0.3.0/tensor_batch_norm.bend as Norm
13import stelliferous@0.0.3.0/tensor_loss.bend as Loss
14import stelliferous@0.0.3.0/tensor_optimizer.bend as Optimizer
15
16def Tensor() -> Type:
17 Tensors.Tensor<F32>
18
19# The gradient of one parameter: none before backward and after sgd_step.
20type Gradient is Type:
21 Empty{}
22 Held{tensor: Tensors.Tensor<F32>}
23
24# A convolution of weights [O,C,k,k] and bias [O]; a batch normalization of
25# gamma, beta and the running statistics, each [C]; an activation; the global
26# average pool of [N,C,H,W] to [N,C,1,1].
27type Layer is Type:
28 Convolution{weights: Tensors.Tensor<F32>,bias: Tensors.Tensor<F32>,stride: U32,padding: U32,weight_gradient: Gradient,bias_gradient: Gradient}
29 BatchNorm{gamma: Tensors.Tensor<F32>,beta: Tensors.Tensor<F32>,running_mean: Tensors.Tensor<F32>,running_variance: Tensors.Tensor<F32>,
30 epsilon: F32,momentum: F32,gamma_gradient: Gradient,beta_gradient: Gradient}
31 Activation{activation: Activations.Activation}
32 AveragePool{}
33
34# What one backward call of a layer reads.
35type Saved is Type:
36 Convolved{saved: ConvLayer.Saved()}
37 Normalized{saved: Norm.Saved()}
38 Activated{input: Tensors.Tensor<F32>}
39 Pooled{shape: List<&2,U32>}
40
41# ---- Layers ----
42
43def conv2d_layer(weights: Tensor(),bias: Tensor(),stride: U32,padding: U32) -> Layer:
44 Convolution{weights,bias,stride,padding,Empty{},Empty{}}
45
46# gamma 1, beta 0, running mean 0 and running variance 1, as in PyTorch.
47def batch_norm_layer_with(+channels: U32,epsilon: F32,momentum: F32) -> Layer:
48 BatchNorm{F.ones([channels]),F.zeros([channels]),F.zeros([channels]),F.ones([channels]),epsilon,momentum,Empty{},Empty{}}
49
50# epsilon 1e-5 and momentum 0.1, as in PyTorch.
51def batch_norm_layer(+channels: U32) -> Layer:
52 batch_norm_layer_with(channels,0.00001,0.1)
53
54# A batch normalization with the given parameters and running statistics.
55def batch_norm_layer_from(gamma: Tensor(),beta: Tensor(),running_mean: Tensor(),running_variance: Tensor(),epsilon: F32,momentum: F32) -> Layer:
56 BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,Empty{},Empty{}}
57
58def activation_layer(activation: Activations.Activation) -> Layer:
59 Activation{activation}
60
61def average_pool_layer() -> Layer:
62 AveragePool{}
63
64# ---- Forward ----
65
66def convolved(+stride: U32,+padding: U32,weight_gradient: Gradient,bias_gradient: Gradient,
67 result: Tensor() & (ConvLayer.Saved() & (Tensor() & Tensor()))) -> Tensor() & (Saved & Layer):
68 (output,(saved,(weights,bias))) = result
69 (output,(Convolved{saved},Convolution{weights,bias,stride,padding,weight_gradient,bias_gradient}))
70
71def normalized(+epsilon: F32,+momentum: F32,gamma_gradient: Gradient,beta_gradient: Gradient,result: Norm.Forward()) -> Tensor() & (Saved & Layer):
72 (output,(saved,((gamma,beta),(running_mean,running_variance)))) = result
73 (output,(Normalized{saved},BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,gamma_gradient,beta_gradient}))
74
75def activated(+activation: Activations.Activation,pair: Tensor() & Tensor()) -> Tensor() & (Saved & Layer):
76 (output,input) = pair
77 (output,(Activated{input},Activation{activation}))
78
79# [N,C,H,W] to [N,C,1,1]; another rank gives the extents unchanged.
80def pooled_extents(shape: List<&2,U32>) -> List<&2,U32>:
81 match shape:
82 case [images,channels,height,width]: [images,channels,1,1]
83 case shape: shape
84
85def pooled_forward(measured: Tensor() & List<&2,U32>) -> Tensor() & (Saved & Layer):
86 (x,+shape) = measured
87 (F.reshape(F.mean_axes(x,[2,3]),pooled_extents(shape)),(Pooled{shape},AveragePool{}))
88
89# (the output, (what backward reads, the layer)). A convolution keeps its
90# input and output; a batch normalization normalizes with the statistics of
91# the batch and moves its running statistics; the pool averages each channel
92# of each image.
93def forward(x: Tensor(),layer: Layer) -> Tensor() & (Saved & Layer):
94 match layer:
95 case Convolution{weights,bias,+stride,+padding,weight_gradient,bias_gradient}:
96 convolved(stride,padding,weight_gradient,bias_gradient,ConvLayer.forward(x,weights,bias,ConvLayer.settings(stride,padding,Activations.identity())))
97 case BatchNorm{gamma,beta,running_mean,running_variance,+epsilon,+momentum,gamma_gradient,beta_gradient}:
98 normalized(epsilon,momentum,gamma_gradient,beta_gradient,Norm.forward(x,(gamma,beta),(running_mean,running_variance),Norm.settings(epsilon,momentum)))
99 case Activation{+activation}: activated(activation,F.activate_forward(x,activation))
100 case AveragePool{}: pooled_forward(Tensors.shape(~F32,x))
101
102# ---- Backward ----
103
104def convolution_back(+stride: U32,+padding: U32,bias: Tensor(),result: Tensor() & (Tensor() & (Tensor() & Tensor()))) -> Tensor() & Layer:
105 (dx,(weight_gradient,(bias_gradient,weights))) = result
106 (dx,Convolution{weights,bias,stride,padding,Held{weight_gradient},Held{bias_gradient}})
107
108def norm_back(beta: Tensor(),running_mean: Tensor(),running_variance: Tensor(),+epsilon: F32,+momentum: F32,result: Norm.Backward()) -> Tensor() & Layer:
109 (dx,(gamma_gradient,(beta_gradient,gamma))) = result
110 (dx,BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,Held{gamma_gradient},Held{beta_gradient}})
111
112# The average of H * W positions spreads its gradient over them.
113def pooled_backward(dy: Tensor(),shape: List<&2,U32>) -> Tensor():
114 match shape:
115 case [images,channels,+height,+width]:
116 F.broadcast_to(F.div_scalar(dy,U32.to_f32((height * width : U32))),[images,channels,height,width])
117 case shape: Tensors.rejected(~F32,"backward",dy)
118
119# A saved value of another kind of layer.
120def unmatched(dy: Tensor(),saved: Saved,layer: Layer) -> Tensor() & Layer:
121 (Tensors.rejected(~F32,"backward",dy),layer)
122
123def convolution_backward(dy: Tensor(),saved: Saved,weights: Tensor(),bias: Tensor(),+stride: U32,+padding: U32,weight_gradient: Gradient,bias_gradient: Gradient) -> Tensor() & Layer:
124 match saved:
125 case Convolved{kept}:
126 convolution_back(stride,padding,bias,ConvLayer.backward(dy,kept,weights,ConvLayer.settings(stride,padding,Activations.identity())))
127 case saved: unmatched(dy,saved,Convolution{weights,bias,stride,padding,weight_gradient,bias_gradient})
128
129def norm_backward(dy: Tensor(),saved: Saved,gamma: Tensor(),beta: Tensor(),running_mean: Tensor(),running_variance: Tensor(),+epsilon: F32,+momentum: F32,
130 gamma_gradient: Gradient,beta_gradient: Gradient) -> Tensor() & Layer:
131 match saved:
132 case Normalized{kept}: norm_back(beta,running_mean,running_variance,epsilon,momentum,Norm.backward(dy,kept,gamma))
133 case saved: unmatched(dy,saved,BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,gamma_gradient,beta_gradient})
134
135def activation_backward(dy: Tensor(),saved: Saved,+activation: Activations.Activation) -> Tensor() & Layer:
136 match saved:
137 case Activated{input}: (F.activate_backward(dy,input,activation),Activation{activation})
138 case saved: unmatched(dy,saved,Activation{activation})
139
140def pool_backward(dy: Tensor(),saved: Saved) -> Tensor() & Layer:
141 match saved:
142 case Pooled{shape}: (pooled_backward(dy,shape),AveragePool{})
143 case saved: unmatched(dy,saved,AveragePool{})
144
145# dy, the gradient of the loss with respect to the layer output, and what
146# forward saved: (the gradient of the layer input, the layer holding the
147# gradients of its parameters). A saved value of another kind of layer gives
148# a failed input gradient.
149def backward(dy: Tensor(),saved: Saved,layer: Layer) -> Tensor() & Layer:
150 match layer:
151 case Convolution{weights,bias,+stride,+padding,weight_gradient,bias_gradient}:
152 convolution_backward(dy,saved,weights,bias,stride,padding,weight_gradient,bias_gradient)
153 case BatchNorm{gamma,beta,running_mean,running_variance,+epsilon,+momentum,gamma_gradient,beta_gradient}:
154 norm_backward(dy,saved,gamma,beta,running_mean,running_variance,epsilon,momentum,gamma_gradient,beta_gradient)
155 case Activation{+activation}: activation_backward(dy,saved,activation)
156 case AveragePool{}: pool_backward(dy,saved)
157
158# ---- Evaluation ----
159
160def convolution_evaluated(x: Tensor(),+stride: U32,+padding: U32,weight_gradient: Gradient,bias_gradient: Gradient,
161 weight_copies: Tensor() & Tensor(),bias_copies: Tensor() & Tensor()) -> Tensor() & Layer:
162 (weights,weights_kept) = weight_copies
163 (bias,bias_kept) = bias_copies
164 (F.conv2d(x,weights,bias,stride,padding,Activations.identity()),Convolution{weights_kept,bias_kept,stride,padding,weight_gradient,bias_gradient})
165
166def norm_evaluated(+epsilon: F32,+momentum: F32,gamma_gradient: Gradient,beta_gradient: Gradient,result: Norm.Evaluated()) -> Tensor() & Layer:
167 (output,((gamma,beta),(running_mean,running_variance))) = result
168 (output,BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,gamma_gradient,beta_gradient})
169
170def pooled_evaluated(measured: Tensor() & List<&2,U32>) -> Tensor() & Layer:
171 (x,shape) = measured
172 (F.reshape(F.mean_axes(x,[2,3]),pooled_extents(shape)),AveragePool{})
173
174# The output in inference mode: a batch normalization uses its running
175# statistics and leaves them unchanged.
176def evaluate(x: Tensor(),layer: Layer) -> Tensor() & Layer:
177 match layer:
178 case Convolution{weights,bias,+stride,+padding,weight_gradient,bias_gradient}:
179 convolution_evaluated(x,stride,padding,weight_gradient,bias_gradient,F.copy(weights),F.copy(bias))
180 case BatchNorm{gamma,beta,running_mean,running_variance,+epsilon,+momentum,gamma_gradient,beta_gradient}:
181 norm_evaluated(epsilon,momentum,gamma_gradient,beta_gradient,Norm.evaluate(x,(gamma,beta),(running_mean,running_variance),Norm.settings(epsilon,momentum)))
182 case Activation{+activation}: (F.activate(x,activation),Activation{activation})
183 case AveragePool{}: pooled_evaluated(Tensors.shape(~F32,x))
184
185# ---- Updates ----
186
187def stepped(parameter: Tensor(),gradient: Gradient,+rate: F32) -> Tensor():
188 match gradient:
189 case Empty{}: parameter
190 case Held{tensor}: Optimizer.sgd(parameter,tensor,rate)
191
192# Every parameter with a gradient becomes parameter - gradient * rate, and the
193# gradients are cleared. A layer without gradients comes back unchanged.
194def sgd_step(layer: Layer,+rate: F32) -> Layer:
195 match layer:
196 case Convolution{weights,bias,stride,padding,weight_gradient,bias_gradient}:
197 Convolution{stepped(weights,weight_gradient,rate),stepped(bias,bias_gradient,rate),stride,padding,Empty{},Empty{}}
198 case BatchNorm{gamma,beta,running_mean,running_variance,epsilon,momentum,gamma_gradient,beta_gradient}:
199 BatchNorm{stepped(gamma,gamma_gradient,rate),stepped(beta,beta_gradient,rate),running_mean,running_variance,epsilon,momentum,Empty{},Empty{}}
200 case Activation{activation}: Activation{activation}
201 case AveragePool{}: AveragePool{}
202
203# ---- Losses ----
204
205# The mean squared error of prediction against a target of its shape:
206# (the loss, (its gradient with respect to prediction, the target)).
207def mse(prediction: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
208 Loss.mse(prediction,target,Loss.mean())
209
210# The mean cross entropy of logits [N,K] against a target [N,K] whose rows
211# are class probabilities: (the loss, (its gradient with respect to logits,
212# the target)).
213def cross_entropy(logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
214 Loss.cross_entropy(logits,target)
215
216# The class numbers [N], each a float, as one-hot rows [N,classes].
217def one_hot(labels: Tensor(),+classes: U32) -> Tensor():
218 Loss.one_hot(labels,classes)