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. 7 import Base 8 import ./math_activation.bend as Activations 9 import ./tensor.bend as Tensors 10 import ./tensor_f32.bend as F 11 import ./tensor_layer.bend as ConvLayer 12 import ./tensor_batch_norm.bend as Norm 13 import ./tensor_loss.bend as Loss 14 import ./tensor_optimizer.bend as Optimizer 15 16 def Tensor() -> Type: 17 Tensors.Tensor 18 19 # The gradient of one parameter: none before backward and after sgd_step. 20 type Gradient is Type: 21 Empty{} 22 Held{tensor: Tensors.Tensor} 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]. 27 type Layer is Type: 28 Convolution{weights: Tensors.Tensor,bias: Tensors.Tensor,stride: U32,padding: U32,weight_gradient: Gradient,bias_gradient: Gradient} 29 BatchNorm{gamma: Tensors.Tensor,beta: Tensors.Tensor,running_mean: Tensors.Tensor,running_variance: Tensors.Tensor, 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. 35 type Saved is Type: 36 Convolved{saved: ConvLayer.Saved()} 37 Normalized{saved: Norm.Saved()} 38 Activated{input: Tensors.Tensor} 39 Pooled{shape: List<&2,U32>} 40 41 # ---- Layers ---- 42 43 def 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. 47 def 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. 51 def 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. 55 def 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 58 def activation_layer(activation: Activations.Activation) -> Layer: 59 Activation{activation} 60 61 def average_pool_layer() -> Layer: 62 AveragePool{} 63 64 # ---- Forward ---- 65 66 def 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 71 def 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 75 def 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. 80 def 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 85 def 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. 93 def 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 104 def 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 108 def 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. 113 def 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. 120 def unmatched(dy: Tensor(),saved: Saved,layer: Layer) -> Tensor() & Layer: 121 (Tensors.rejected(~F32,"backward",dy),layer) 122 123 def 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 129 def 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 135 def 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 140 def 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. 149 def 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 160 def 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 166 def 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 170 def 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. 176 def 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 187 def 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. 194 def 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)). 207 def 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)). 213 def 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]. 217 def one_hot(labels: Tensor(),+classes: U32) -> Tensor(): 218 Loss.one_hot(labels,classes)