tensor_convolution_gradient.bend
stelliferous@0.0.3.0 · 106 lines · 0 definitions
1# Gradients of conv2d in a fixed FP32 order. Every function gives the
2# tensors it only reads back to its caller.
3import Base
4import stelliferous@0.0.3.0/tensor.bend as Tensors
5import stelliferous@0.0.3.0/tensor_band.bend as Bands
6import stelliferous@0.0.3.0/tensor_convolution.bend as TensorConvolution
7import stelliferous@0.0.3.0/kernel_input_gradient.bend as InputKernel
8import stelliferous@0.0.3.0/kernel_shape.bend as KernelShape
9import stelliferous@0.0.3.0/storage_buffer.bend as Storage
10import stelliferous@0.0.3.0/default_machine_limits.bend as Default
11import stelliferous@0.0.3.0/machine_limits.bend as Machine
12import stelliferous@0.0.3.0/tensor_parameter_gradient.bend as ParameterGradient
13
14def Tensor() -> Type:
15 Tensors.Tensor<F32>
16
17type Arguments is Data:
18 Arguments{geometry: InputKernel.Geometry,batch: U32,batched: Bool}
19
20def input_dimensions(arguments: Arguments) -> KernelShape.Dimensions:
21 Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
22 KernelShape.Dimensions{U32.to_nat(c),U32.to_nat(o),U32.to_nat(h),U32.to_nat(w),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p)}
23
24# The bands of the input gradient: rows of every image of the batch.
25def input_band_geometry(arguments: Arguments) -> Bands.Geometry:
26 Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
27 Bands.Geometry{batch,c,h,w}
28
29def input_shape(arguments: Arguments) -> List<&2,U32>:
30 Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
31 TensorConvolution.output_shape(batched,batch,c,h,w)
32
33# The limits are an argument so that a proof can keep them symbolic.
34def input_valid_with(+arguments: Arguments,d_length: Nat,+weight_length: Nat,+limits: Machine.Limits) -> Bool:
35 Arguments{+geometry,+batch,batched} = arguments
36 InputKernel.Geometry{+c,+o,+h,+w,+k,+s,+p,+gh,+gw} = geometry
37 +dimensions = input_dimensions(arguments)
38 +input_length = Nat.mul(U32.to_nat(c),Nat.mul(U32.to_nat(h),U32.to_nat(w)))
39 KernelShape.valid(dimensions,input_length,weight_length,U32.to_nat(o)) &&
40 KernelShape.supports_direct(limits,dimensions,input_length,weight_length) &&
41 Nat.is_eq(U32.to_nat(gh),KernelShape.axis_size(U32.to_nat(h),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p))) &&
42 Nat.is_eq(U32.to_nat(gw),KernelShape.axis_size(U32.to_nat(w),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p))) &&
43 Nat.is_eq(d_length,Nat.mul(U32.to_nat(batch),Nat.mul(U32.to_nat(o),Nat.mul(U32.to_nat(gh),U32.to_nat(gw))))) &&
44 InputKernel.fits(geometry,Machine.capacity(limits))
45
46def input_valid(arguments: Arguments,d_length: Nat,weight_length: Nat) -> Bool:
47 input_valid_with(arguments,d_length,weight_length,Default.limits())
48
49def input_result(+arguments: Arguments,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32,result: InputKernel.Outcome) -> Tensor() & Tensor():
50 InputKernel.Outcome{weights,output} = result
51 (Tensors.Tensor{output,input_band_geometry(arguments),input_shape(arguments),0.0,Tensors.Usable{},workers},Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
52
53def input_checked(accepted: Bool,+arguments: Arguments,+count: U32,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
54 match accepted:
55 case False{}:
56 (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidArgument{input_shape(arguments)},input_shape(arguments),0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
57 case True{}:
58 Arguments{+geometry,batch,batched} = arguments
59 input_result(arguments,workers,weight_shape,weight_fill,weight_workers,InputKernel.run(geometry,U32.to_nat(count),U32.max(workers,U32.max(weight_workers,1)),gradient,weights))
60
61def input_count(count: Maybe<&2,U32>,+arguments: Arguments,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
62 match count:
63 case None{}:
64 (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidShape{input_shape(arguments)},input_shape(arguments),0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
65 case Some{+count}:
66 Storage.Buffer{d,+d_length,d_certificate} = gradient
67 Storage.Buffer{w,+weight_length,w_certificate} = weights
68 input_checked(input_valid(arguments,d_length,weight_length),arguments,count,gradient,weights,workers,weight_shape,weight_fill,weight_workers)
69
70def input_geometry(+stride: U32,+padding: U32,+height: U32,+width: U32,shape: Maybe<&2,TensorConvolution.InputShape>,+weight_shape: List<&2,U32>,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
71 match shape weight_shape:
72 case Some{TensorConvolution.InputShape{+batch,+batched,+outputs,+gh,+gw}} Con{wo,Con{+channels,Con{+kernel,Con{kw,Nil{}}}}}:
73 +arguments = {Arguments{InputKernel.Geometry{channels,outputs,height,width,kernel,stride,padding,gh,gw},batch,batched} : Arguments}
74 input_count(Bool.pick(Maybe<&2,U32>,U32.is_eq(outputs,wo) && U32.is_eq(kernel,kw),Tensors.count_of(input_shape(arguments)),None{}),arguments,gradient,weights,workers,weight_shape,weight_fill,weight_workers)
75 case shape weight_shape:
76 (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidShape{weight_shape},weight_shape,0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
77
78def input_status(status: Tensors.Status,+stride: U32,+padding: U32,+height: U32,+width: U32,gradient: Storage.Buffer<F32>,+d_shape: List<&2,U32>,+d_fill: F32,+workers: U32,weights: Storage.Buffer<F32>,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_status: Tensors.Status,+weight_workers: U32) -> Tensor() & Tensor():
79 match status:
80 case Tensors.Failed{problem}:
81 (Tensors.whole(~F32,gradient,d_shape,d_fill,Tensors.Failed{problem},workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,weight_status,weight_workers))
82 case Tensors.Usable{}:
83 input_geometry(stride,padding,height,width,TensorConvolution.input_shape(d_shape),weight_shape,gradient,weights,workers,weight_fill,weight_workers)
84
85# d [N,O,Ho,Wo] or [O,Ho,Wo] and weights [O,C,k,k] of a convolution of an
86# input of the given height and width: (the gradient of that input, the
87# weights). Several input sizes give the same output size when the stride is
88# above 1, so height and width are given.
89def input(d: Tensor(),weights: Tensor(),+stride: U32,+padding: U32,+height: U32,+width: U32) -> Tensor() & Tensor():
90 Tensors.with_whole(~F32,~(Tensor() & Tensor()),d,gradient => d_shape => d_fill => d_status => workers =>
91 Tensors.with_whole(~F32,~(Tensor() & Tensor()),weights,weights => weight_shape => weight_fill => weight_status => weight_workers =>
92 +weight_status = weight_status
93 input_status(Tensors.first_failure(d_status,weight_status),stride,padding,height,width,gradient,d_shape,d_fill,workers,weights,weight_shape,weight_fill,weight_status,weight_workers)))
94
95# The group length of the weight and bias gradient contract.
96def group_length() -> U32:
97 256
98
99# x [N,C,H,W], d [N,O,Ho,Wo], or both without their batch axes.
100# Returns (weight gradient [O,C,k,k], (x,d)); all sources are read only.
101def weight(x: Tensor(),d: Tensor(),kernel: U32,stride: U32,padding: U32) -> Tensor() & (Tensor() & Tensor()):
102 ParameterGradient.weight(x,d,kernel,stride,padding,group_length())
103
104# d [N,O,Ho,Wo] or [O,Ho,Wo]: (bias gradient [O], d).
105def bias(d: Tensor()) -> Tensor() & Tensor():
106 ParameterGradient.bias(d,group_length())