tensor_layer.bend
stelliferous@0.0.3.0 · 103 lines · 0 definitions
1# A convolution followed by an activation, as a training layer: forward keeps
2# what backward reads, backward gives the gradients in a fixed FP32 order.
3import Base
4import stelliferous@0.0.3.0/math_activation.bend as Activations
5import stelliferous@0.0.3.0/tensor.bend as Tensors
6import stelliferous@0.0.3.0/tensor_f32.bend as F
7import stelliferous@0.0.3.0/tensor_convolution_gradient.bend as Gradient
8
9def Tensor() -> Type:
10 Tensors.Tensor<F32>
11
12type Settings is Data:
13 Settings{stride: U32,padding: U32,activation: Activations.Activation}
14
15def settings(stride: U32,padding: U32,activation: Activations.Activation) -> Settings:
16 Settings{stride,padding,activation}
17
18# The layer input and the convolution result before the activation. One
19# backward call uses it.
20def Saved() -> Type:
21 Tensor() & Tensor()
22
23def forward_activated(x: Tensor(),weights: Tensor(),bias: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Saved() & (Tensor() & Tensor())):
24 (output,z) = pair
25 (output,((x,z),(weights,bias)))
26
27def forward_copied(settings: Settings,xs: Tensor() & Tensor(),ws: Tensor() & Tensor(),bs: Tensor() & Tensor()) -> Tensor() & (Saved() & (Tensor() & Tensor())):
28 Settings{stride,padding,+activation} = settings
29 (x,x_kept) = xs
30 (weights,weights_kept) = ws
31 (bias,bias_kept) = bs
32 forward_activated(x_kept,weights_kept,bias_kept,F.activate_forward(F.conv2d(x,weights,bias,stride,padding,Activations.identity()),activation))
33
34# x [N,C,H,W] or [C,H,W], weights [O,C,k,k] and bias [O]:
35# (the activated output, (what backward reads, (weights, bias))).
36def forward(x: Tensor(),weights: Tensor(),bias: Tensor(),settings: Settings) -> Tensor() & (Saved() & (Tensor() & Tensor())):
37 forward_copied(settings,F.copy(x),F.copy(weights),F.copy(bias))
38
39def kernel(shape: List<&2,U32>) -> U32:
40 match shape:
41 case [outputs,channels,kernel,width]: kernel
42 case shape: 0
43
44# The gradients of the parameters with what the input gradient reads:
45# (weight gradient, (bias gradient, (dz, (weights, the shape of x)))).
46def Parameters() -> Type:
47 Tensor() & (Tensor() & (Tensor() & (Tensor() & List<&2,U32>)))
48
49def biased(weight_gradient: Tensor(),weights: Tensor(),shape: List<&2,U32>,pair: Tensor() & Tensor()) -> Parameters():
50 (bias_gradient,dz) = pair
51 (weight_gradient,(bias_gradient,(dz,(weights,shape))))
52
53def measured(weight_gradient: Tensor(),weights: Tensor(),dz: Tensor(),input: Tensor() & List<&2,U32>) -> Parameters():
54 (x,shape) = input
55 biased(weight_gradient,weights,shape,Gradient.bias(dz))
56
57def weighted(weights: Tensor(),result: Tensor() & (Tensor() & Tensor())) -> Parameters():
58 (weight_gradient,(x,dz)) = result
59 measured(weight_gradient,weights,dz,Tensors.shape(~F32,x))
60
61def shaped(stride: U32,padding: U32,x: Tensor(),dz: Tensor(),measured: Tensor() & List<&2,U32>) -> Parameters():
62 (weights,shape) = measured
63 weighted(weights,Gradient.weight(x,dz,kernel(shape),stride,padding))
64
65# dz = dy * activation'(z), then the gradients of the weights and the bias.
66def parameters(dy: Tensor(),saved: Saved(),weights: Tensor(),settings: Settings) -> Parameters():
67 (x,z) = saved
68 Settings{stride,padding,+activation} = settings
69 shaped(stride,padding,x,F.activate_backward(dy,z,activation),Tensors.shape(~F32,weights))
70
71def finished(weight_gradient: Tensor(),bias_gradient: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
72 (dx,weights) = pair
73 (dx,(weight_gradient,(bias_gradient,weights)))
74
75# The height and width of [N,C,H,W] or [C,H,W]; zero for another rank.
76def image(shape: List<&2,U32>) -> U32 & U32:
77 match shape:
78 case [images,channels,height,width]: (height,width)
79 case [channels,height,width]: (height,width)
80 case shape: (0,0)
81
82def sized(+stride: U32,+padding: U32,weight_gradient: Tensor(),bias_gradient: Tensor(),dz: Tensor(),weights: Tensor(),size: U32 & U32) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
83 (height,width) = size
84 finished(weight_gradient,bias_gradient,Gradient.input(dz,weights,stride,padding,height,width))
85
86def propagated(+stride: U32,+padding: U32,result: Parameters()) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
87 (weight_gradient,(bias_gradient,(dz,(weights,shape)))) = result
88 sized(stride,padding,weight_gradient,bias_gradient,dz,weights,image(shape))
89
90# dy, the gradient of the loss with respect to the layer output:
91# (the gradient of the layer input, (weight gradient, (bias gradient, weights))).
92def backward(dy: Tensor(),saved: Saved(),weights: Tensor(),+settings: Settings) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
93 Settings{+stride,+padding,activation} = settings
94 propagated(stride,padding,parameters(dy,saved,weights,settings))
95
96def without_input(result: Parameters()) -> Tensor() & (Tensor() & Tensor()):
97 (weight_gradient,(bias_gradient,(dz,(weights,shape)))) = result
98 (weight_gradient,(bias_gradient,weights))
99
100# backward for a layer whose input needs no gradient:
101# (weight gradient, (bias gradient, weights)).
102def backward_parameters(dy: Tensor(),saved: Saved(),weights: Tensor(),settings: Settings) -> Tensor() & (Tensor() & Tensor()):
103 without_input(parameters(dy,saved,weights,settings))