tensor_batch_norm.bend
stelliferous@0.0.3.0 · 179 lines · 0 definitions
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.
6import Base
7import stelliferous@0.0.3.0/tensor.bend as Tensors
8import stelliferous@0.0.3.0/tensor_f32.bend as F
9
10def Tensor() -> Type:
11 Tensors.Tensor<F32>
12
13type Settings is Data:
14 Settings{epsilon: F32,momentum: F32}
15
16def settings(epsilon: F32,momentum: F32) -> Settings:
17 Settings{epsilon,momentum}
18
19# gamma and beta, each [C].
20def Parameters() -> Type:
21 Tensor() & Tensor()
22
23# The running mean and the running variance, each [C].
24def Running() -> Type:
25 Tensor() & Tensor()
26
27# The normalized input and 1 / sqrt(variance + epsilon). One backward call
28# uses it.
29def Saved() -> Type:
30 Tensor() & Tensor()
31
32def 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.
38def 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].
44def spread(tensor: Tensor(),+channels: U32) -> Tensor():
45 F.reshape(tensor,[1,channels,1,1])
46
47# 1 / sqrt(variance + epsilon).
48def 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.
52def 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.
56def 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))).
64def Forward() -> Type:
65 Tensor() & (Saved() & (Parameters() & Running()))
66
67def 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
71def 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
78def 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
83def 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
90def 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
95def 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
100def 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
104def 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
108def 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).
116def forward(x: Tensor(),parameters: Parameters(),running: Running(),+settings: Settings) -> Forward():
117 shaped(settings,F.shape(x),parameters,running)
118
119# (dx, (dgamma, (dbeta, gamma))).
120def Backward() -> Type:
121 Tensor() & (Tensor() & (Tensor() & Tensor()))
122
123def 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
131def 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
135def 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
141def 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).
150def 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)).
156def Evaluated() -> Type:
157 Tensor() & (Parameters() & Running())
158
159def evaluation(mean: Tensor(),variance: Tensor(),result: Tensor() & Parameters()) -> Evaluated():
160 (output,parameters) = result
161 (output,(parameters,(mean,variance)))
162
163def 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
170def 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.
178def evaluate(x: Tensor(),parameters: Parameters(),running: Running(),+settings: Settings) -> Evaluated():
179 evaluate_shaped(settings,F.shape(x),parameters,running)