1 # Batch normalization of [N,C,H,W] over the images and positions of each 2 # channel, as a training layer: forward normalizes with the statistics of the 3 # batch and moves the running statistics, backward gives the gradients, 4 # evaluate normalizes with the running statistics. Every step is a public 5 # operation in a fixed FP32 order; the sums are those of sum_axes. 6 import Base 7 import ./tensor.bend as Tensors 8 import ./tensor_f32.bend as F 9 10 def Tensor() -> Type: 11 Tensors.Tensor 12 13 type Settings is Data: 14 Settings{epsilon: F32,momentum: F32} 15 16 def settings(epsilon: F32,momentum: F32) -> Settings: 17 Settings{epsilon,momentum} 18 19 # gamma and beta, each [C]. 20 def Parameters() -> Type: 21 Tensor() & Tensor() 22 23 # The running mean and the running variance, each [C]. 24 def Running() -> Type: 25 Tensor() & Tensor() 26 27 # The normalized input and 1 / sqrt(variance + epsilon). One backward call 28 # uses it. 29 def Saved() -> Type: 30 Tensor() & Tensor() 31 32 def channels(shape: List<&2,U32>) -> U32: 33 match shape: 34 case [images,channels,height,width]: channels 35 case shape: 0 36 37 # The number of values of each channel, N * H * W. 38 def count(shape: List<&2,U32>) -> U32: 39 match shape: 40 case [images,channels,height,width]: (images * (height * width : U32) : U32) 41 case shape: 0 42 43 # A per-channel tensor [C] seen as [1,C,1,1], to broadcast over [N,C,H,W]. 44 def spread(tensor: Tensor(),+channels: U32) -> Tensor(): 45 F.reshape(tensor,[1,channels,1,1]) 46 47 # 1 / sqrt(variance + epsilon). 48 def inverse_deviation(variance: Tensor(),+channels: U32,+epsilon: F32) -> Tensor(): 49 F.div(F.ones([channels]),F.sqrt(F.add_scalar(variance,epsilon))) 50 51 # running * (1 - momentum) + statistic * momentum. 52 def moved(running: Tensor(),statistic: Tensor(),+momentum: F32) -> Tensor(): 53 F.add(F.mul_scalar(running,(1.0 - momentum : F32)),F.mul_scalar(statistic,momentum)) 54 55 # xhat * gamma + beta, and the parameters back. 56 def affine(xhat: Tensor(),+channels: U32,gammas: Tensor() & Tensor(),betas: Tensor() & Tensor()) -> Tensor() & Parameters(): 57 (gamma,gamma_kept) = gammas 58 (beta,beta_kept) = betas 59 (F.add(F.mul(xhat,spread(gamma,channels)),spread(beta,channels)),(gamma_kept,beta_kept)) 60 61 # ---- Training ---- 62 63 # (the output, (what backward reads, (parameters, running statistics))). 64 def Forward() -> Type: 65 Tensor() & (Saved() & (Parameters() & Running())) 66 67 def finished(xhat: Tensor(),deviation: Tensor(),running_mean: Tensor(),running_variance: Tensor(),result: Tensor() & Parameters()) -> Forward(): 68 (output,parameters) = result 69 (output,((xhat,deviation),(parameters,(running_mean,running_variance)))) 70 71 def normalized(+channels: U32,+momentum: F32,mean: Tensor(),unbiased: Tensor(),deviation: Tensor(),xhats: Tensor() & Tensor(),parameters: Parameters(), 72 running: Running()) -> Forward(): 73 (xhat,xhat_kept) = xhats 74 (gamma,beta) = parameters 75 (running_mean,running_variance) = running 76 finished(xhat_kept,deviation,moved(running_mean,mean,momentum),moved(running_variance,unbiased,momentum),affine(xhat,channels,F.copy(gamma),F.copy(beta))) 77 78 def scaled(+channels: U32,+momentum: F32,mean: Tensor(),difference: Tensor(),unbiased: Tensor(),deviations: Tensor() & Tensor(),parameters: Parameters(), 79 running: Running()) -> Forward(): 80 (deviation,deviation_kept) = deviations 81 normalized(channels,momentum,mean,unbiased,deviation_kept,F.copy(F.mul(difference,spread(deviation,channels))),parameters,running) 82 83 def deviated(+settings: Settings,+channels: U32,+values: U32,mean: Tensor(),difference: Tensor(),sums: Tensor() & Tensor(),parameters: Parameters(), 84 running: Running()) -> Forward(): 85 Settings{+epsilon,+momentum} = settings 86 (squares,squares_kept) = sums 87 scaled(channels,momentum,mean,difference,F.div_scalar(squares_kept,U32.to_f32((values - 1 : U32))), 88 F.copy(inverse_deviation(F.div_scalar(squares,U32.to_f32(values)),channels,epsilon)),parameters,running) 89 90 def squared(+settings: Settings,+channels: U32,+values: U32,mean: Tensor(),difference: Tensor(),pair: Tensor() & Tensor(),parameters: Parameters(), 91 running: Running()) -> Forward(): 92 (left,right) = pair 93 deviated(settings,channels,values,mean,difference,F.copy(F.sum_axes(F.mul(left,right),[0,2,3])),parameters,running) 94 95 def differenced(+settings: Settings,+channels: U32,+values: U32,mean: Tensor(),differences: Tensor() & Tensor(),parameters: Parameters(), 96 running: Running()) -> Forward(): 97 (difference,difference_kept) = differences 98 squared(settings,channels,values,mean,difference_kept,F.copy(difference),parameters,running) 99 100 def centered(+settings: Settings,+channels: U32,+values: U32,means: Tensor() & Tensor(),x: Tensor(),parameters: Parameters(),running: Running()) -> Forward(): 101 (mean,mean_kept) = means 102 differenced(settings,channels,values,mean_kept,F.copy(F.sub(x,spread(mean,channels))),parameters,running) 103 104 def split(+settings: Settings,+channels: U32,+values: U32,xs: Tensor() & Tensor(),parameters: Parameters(),running: Running()) -> Forward(): 105 (x,x_kept) = xs 106 centered(settings,channels,values,F.copy(F.div_scalar(F.sum_axes(x,[0,2,3]),U32.to_f32(values))),x_kept,parameters,running) 107 108 def shaped(+settings: Settings,measured: Tensor() & List<&2,U32>,parameters: Parameters(),running: Running()) -> Forward(): 109 (x,+shape) = measured 110 split(settings,channels(shape),count(shape),F.copy(x),parameters,running) 111 112 # x [N,C,H,W], with M = N * H * W, mean = sum_axes(x,[0,2,3]) / M, d = x - mean 113 # and s = sum_axes(d * d,[0,2,3]): xhat = d * (1 / sqrt(s / M + epsilon)), the 114 # output xhat * gamma + beta, the running mean moved towards mean and the 115 # running variance towards s / (M - 1). 116 def forward(x: Tensor(),parameters: Parameters(),running: Running(),+settings: Settings) -> Forward(): 117 shaped(settings,F.shape(x),parameters,running) 118 119 # (dx, (dgamma, (dbeta, gamma))). 120 def Backward() -> Type: 121 Tensor() & (Tensor() & (Tensor() & Tensor())) 122 123 def propagated(+channels: U32,+values: U32,dy: Tensor(),xhat: Tensor(),scale: Tensor(),dbetas: Tensor() & Tensor(),dgammas: Tensor() & Tensor(), 124 gamma: Tensor()) -> Backward(): 125 (dbeta,dbeta_kept) = dbetas 126 (dgamma,dgamma_kept) = dgammas 127 +total = U32.to_f32(values) 128 (F.mul(F.sub(F.sub(F.mul_scalar(dy,total),spread(dbeta,channels)),F.mul(xhat,spread(dgamma,channels))),spread(F.div_scalar(scale,total),channels)), 129 (dgamma_kept,(dbeta_kept,gamma))) 130 131 def sums(+channels: U32,+values: U32,dys: Tensor() & Tensor(),dy: Tensor(),xhat: Tensor(),xhat_kept: Tensor(),scale: Tensor(),gamma: Tensor()) -> Backward(): 132 (for_beta,for_gamma) = dys 133 propagated(channels,values,dy,xhat_kept,scale,F.copy(F.sum_axes(for_beta,[0,2,3])),F.copy(F.sum_axes(F.mul(for_gamma,xhat),[0,2,3])),gamma) 134 135 def copied(+channels: U32,+values: U32,dys: Tensor() & Tensor(),xhats: Tensor() & Tensor(),deviation: Tensor(),gammas: Tensor() & Tensor()) -> Backward(): 136 (dy,dy_kept) = dys 137 (xhat,xhat_kept) = xhats 138 (gamma,gamma_kept) = gammas 139 sums(channels,values,F.copy(dy),dy_kept,xhat,xhat_kept,F.mul(gamma,deviation),gamma_kept) 140 141 def backward_shaped(measured: Tensor() & List<&2,U32>,saved: Saved(),gamma: Tensor()) -> Backward(): 142 (dy,+shape) = measured 143 (xhat,deviation) = saved 144 copied(channels(shape),count(shape),F.copy(dy),F.copy(xhat),deviation,F.copy(gamma)) 145 146 # dy is the gradient of the loss with respect to the output. With 147 # dbeta = sum_axes(dy,[0,2,3]) and dgamma = sum_axes(dy * xhat,[0,2,3]): 148 # dx = ((dy * M - dbeta) - xhat * dgamma) * ((gamma * deviation) / M), where 149 # deviation is the saved 1 / sqrt(variance + epsilon). 150 def backward(dy: Tensor(),saved: Saved(),gamma: Tensor()) -> Backward(): 151 backward_shaped(F.shape(dy),saved,gamma) 152 153 # ---- Evaluation ---- 154 155 # (the output, (parameters, running statistics)). 156 def Evaluated() -> Type: 157 Tensor() & (Parameters() & Running()) 158 159 def evaluation(mean: Tensor(),variance: Tensor(),result: Tensor() & Parameters()) -> Evaluated(): 160 (output,parameters) = result 161 (output,(parameters,(mean,variance))) 162 163 def evaluated(+channels: U32,+epsilon: F32,x: Tensor(),parameters: Parameters(),means: Tensor() & Tensor(),variances: Tensor() & Tensor()) -> Evaluated(): 164 (gamma,beta) = parameters 165 (mean,mean_kept) = means 166 (variance,variance_kept) = variances 167 evaluation(mean_kept,variance_kept, 168 affine(F.mul(F.sub(x,spread(mean,channels)),spread(inverse_deviation(variance,channels,epsilon),channels)),channels,F.copy(gamma),F.copy(beta))) 169 170 def evaluate_shaped(+settings: Settings,measured: Tensor() & List<&2,U32>,parameters: Parameters(),running: Running()) -> Evaluated(): 171 Settings{+epsilon,momentum} = settings 172 (x,+shape) = measured 173 (running_mean,running_variance) = running 174 evaluated(channels(shape),epsilon,x,parameters,F.copy(running_mean),F.copy(running_variance)) 175 176 # The output with the running statistics: 177 # (x - running mean) * (1 / sqrt(running variance + epsilon)) * gamma + beta. 178 def evaluate(x: Tensor(),parameters: Parameters(),running: Running(),+settings: Settings) -> Evaluated(): 179 evaluate_shaped(settings,F.shape(x),parameters,running)