tensor_loss.bend
stelliferous@0.0.3.0 · 134 lines · 0 definitions
1# Losses with their gradient, in a fixed FP32 order. The prediction becomes
2# the gradient, and the target returns.
3import Base
4import stelliferous@0.0.3.0/tensor.bend as Tensors
5import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
6import stelliferous@0.0.3.0/tensor_layout.bend as Layout
7import stelliferous@0.0.3.0/tensor_operations.bend as Operations
8import stelliferous@0.0.3.0/math_sum.bend as Sum
9import stelliferous@0.0.3.0/tensor_f32.bend as F
10import stelliferous@0.0.3.0/math_fp32.bend as Arithmetic
11
12def Tensor() -> Type:
13 Tensors.Tensor<F32>
14
15type Reduction is Data:
16 Mean{}
17 Sum{}
18
19def mean() -> Reduction:
20 Mean{}
21
22def sum() -> Reduction:
23 Sum{}
24
25def square(value: F32) -> F32:
26 +value = value
27 (value * value : F32)
28
29def reduced(reduction: Reduction,+count: F32,total: Tensor()) -> Tensor():
30 match reduction:
31 case Mean{}: Tensors.combine_scalar(~F32,~F32.div,1.0,total,count)
32 case Sum{}: total
33
34def mse_squares(+reduction: Reduction,+count: F32,target: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
35 (difference,squares) = pair
36 (reduced(reduction,count,Layout.reduce_all(~F32,~Sum.Lanes,~Sum.step,~Sum.total,"mse_loss",6.0,Sum.zero(),
37 Tensors.apply(~F32,~Operations.mapped(~F32,~square),1.0,squares))),
38 (reduced(reduction,count,Tensors.combine_scalar(~F32,~F32.mul,1.0,difference,2.0)),target))
39
40def mse_counted(+reduction: Reduction,target: Tensor(),shaped: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
41 (difference,+shape) = shaped
42 mse_squares(reduction,U32.to_f32(Tensors.count_or_zero(shape)),target,Layout.copy(~F32,difference))
43
44def mse_difference(same: Bool,prediction: Tensor(),kept: Tensor()) -> Tensor():
45 match same:
46 case True{}: Tensors.combine(~F32,~F32.sub,"mse_loss",1.0,prediction,kept)
47 case False{}: Tensors.rejected(~F32,"mse_loss",prediction)
48
49def mse_kept(+reduction: Reduction,same: Bool,prediction: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
50 (target,kept) = pair
51 mse_counted(reduction,target,Tensors.shape(~F32,mse_difference(same,prediction,kept)))
52
53def mse_shaped(+reduction: Reduction,shaped: Tensor() & List<&2,U32>,measured: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
54 (prediction,shape) = shaped
55 (target,target_shape) = measured
56 mse_kept(reduction,Shapes.same_extents(shape,target_shape),prediction,Layout.copy(~F32,target))
57
58# The loss of prediction against a target of the same shape, the gradient of
59# the loss with respect to prediction, and the target. With d = prediction -
60# target: Mean gives sum(d * d) / count and (d * 2) / count, Sum gives
61# sum(d * d) and d * 2; the squares are added in row-major order from +0.
62def mse(prediction: Tensor(),target: Tensor(),+reduction: Reduction) -> Tensor() & (Tensor() & Tensor()):
63 mse_shaped(reduction,Tensors.shape(~F32,prediction),Tensors.shape(~F32,target))
64
65# ---- Cross entropy ----
66
67# The class numbers [N] as rows [N,K] of a target: 1 where the column equals
68# the class number, 0 elsewhere.
69def one_hot_rows(shape: List<&2,U32>,labels: Tensor(),+classes: U32) -> Tensor():
70 match shape:
71 case [images]: F.equal(F.reshape(labels,[images,1]),F.reshape(F.arange(0.0,classes),[1,classes]))
72 case shape: Tensors.rejected(~F32,"one_hot",labels)
73
74def one_hot_measured(measured: Tensor() & List<&2,U32>,+classes: U32) -> Tensor():
75 (labels,shape) = measured
76 one_hot_rows(shape,labels,classes)
77
78def one_hot(labels: Tensor(),+classes: U32) -> Tensor():
79 one_hot_measured(Tensors.shape(~F32,labels),classes)
80
81# The exponentials of a row added in the lanes of math_sum.bend, and the
82# logarithm of their sum.
83def exponential_step(lanes: Sum.Lanes,value: F32) -> Sum.Lanes:
84 Sum.step(lanes,Arithmetic.exp(value))
85
86def logarithm_total(lanes: Sum.Lanes) -> F32:
87 Arithmetic.log(Sum.total(lanes))
88
89def entropy_gradient(+images: U32,log_probabilities: Tensor() & Tensor(),targets: Tensor() & Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
90 (for_loss,for_gradient) = log_probabilities
91 (for_product,for_difference) = targets
92 (F.div_scalar(F.sum(F.mul(for_product,for_loss)),F32.neg(U32.to_f32(images))),
93 (F.div_scalar(F.sub(F.exp(for_gradient),for_difference),U32.to_f32(images)),target))
94
95def entropy_shifted(+images: U32,shifted: Tensor() & Tensor(),targets: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
96 (values,for_sum) = shifted
97 (target,target_kept) = targets
98 entropy_gradient(images,
99 F.copy(F.sub(values,F.reshape(Layout.reduce_axes(~F32,~Sum.Lanes,~exponential_step,~logarithm_total,"cross_entropy",10.5,Sum.zero(),for_sum,[1]),[images,1]))),
100 F.copy(target),target_kept)
101
102def entropy_rows(+images: U32,logits: Tensor() & Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
103 (values,for_maximum) = logits
104 entropy_shifted(images,
105 F.copy(F.sub(values,F.reshape(Layout.reduce_axes(~F32,~F32,~F32.max,~Tensors.same(~F32),"cross_entropy",1.0,F32.div(F32.neg(1.0),0.0),for_maximum,[1]),[images,1]))),
106 F.copy(target))
107
108def entropy_failed(shape: List<&2,U32>,logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
109 (Tensors.failed(~F32,"cross_entropy",Tensors.InvalidShape{shape},[],0.0,1),(Tensors.rejected(~F32,"cross_entropy",logits),target))
110
111def entropy_ranked(shape: List<&2,U32>,logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
112 match shape:
113 case [images,classes]: entropy_rows(images,F.copy(logits),target)
114 case shape: entropy_failed(shape,logits,target)
115
116def entropy_shaped(same: Bool,logits: Tensor(),shape: List<&2,U32>,target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
117 match same:
118 case True{}: entropy_ranked(shape,logits,target)
119 case False{}: entropy_failed(shape,logits,target)
120
121def entropy_measured(measured_logits: Tensor() & List<&2,U32>,measured_target: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
122 (logits,+shape) = measured_logits
123 (target,target_shape) = measured_target
124 entropy_shaped(Shapes.same_extents(shape,target_shape),logits,shape,target)
125
126# logits [N,K] and a target [N,K] whose rows are class probabilities, such as
127# one_hot of the class numbers: the mean loss, its gradient with respect to
128# logits, and the target. With m the maximum of a row and the log
129# probabilities p = (logits - m) - log(sum of exp(logits - m) over the row),
130# the sum over the row in the lanes of math_sum.bend: the loss is
131# sum(target * p) / -N, the sum of math_sum.bend in row-major order, and the
132# gradient (exp(p) - target) / N.
133def cross_entropy(logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
134 entropy_measured(Tensors.shape(~F32,logits),Tensors.shape(~F32,target))