1 module dopt.core.grads.nnet;
2 
3 import dopt.core.grads;
4 import dopt.core.ops;
5 
6 package
7 {
8     void initialize()
9     {
10         import std.functional : toDelegate;
11         
12         registerGradient("convolution", toDelegate(&convolutionGrad));
13         registerGradient("convolutionFeaturesGrad", toDelegate(&convolutionFeaturesGradGrad));
14         registerGradient("maxpool", toDelegate(&maxpoolGrad));
15         registerGradient("softmax", toDelegate(&softmaxGrad));
16         registerGradient("relu", toDelegate(&reluGrad));
17         registerGradient("addBias", toDelegate(&addBiasGrad));
18         registerGradient("batchNormTrain", toDelegate(&batchNormGrad));
19     }
20 }
21 
22 private
23 {
24     Operation[] convolutionGrad(Operation op, Operation parentGrad)
25     {
26         auto padding = op.attributes["padding"].get!(size_t[]);
27         auto stride = op.attributes["stride"].get!(size_t[]);
28 
29         return [
30             convolutionFeaturesGrad(parentGrad, op.deps[1], op.deps[0].shape, padding, stride),
31             convolutionFiltersGrad(parentGrad, op.deps[0], op.deps[1].shape, padding, stride)
32         ];
33     }
34 
35     Operation[] convolutionFeaturesGradGrad(Operation op, Operation parentGrad)
36     {
37         auto padding = op.attributes["padding"].get!(size_t[]);
38         auto stride = op.attributes["stride"].get!(size_t[]);
39 
40         return [
41             convolution(parentGrad, op.deps[1], padding, stride),
42             convolutionFiltersGrad(op.deps[0], parentGrad, op.deps[1].shape, padding, stride)
43         ];
44     }
45 
46     Operation[] maxpoolGrad(Operation op, Operation parentGrad)
47     {
48         return [dopt.core.ops.nnet.maxpoolGrad(parentGrad, op)];
49     }
50 
51     Operation[] softmaxGrad(Operation op, Operation parentGrad)
52     {
53         return [dopt.core.ops.nnet.softmaxGrad(parentGrad, op)];
54     }
55 
56     Operation[] reluGrad(Operation op, Operation parentGrad)
57     {
58         return [dopt.core.ops.nnet.reluGrad(parentGrad, op)];
59     }
60 
61     Operation[] addBiasGrad(Operation op, Operation parentGrad)
62     {
63         return [parentGrad, dopt.core.ops.nnet.addBiasGrad(parentGrad)];
64     }
65 
66     Operation[] batchNormGrad(Operation op, Operation parentGrad)
67     {
68         auto trimmedGrad = parentGrad.slice([0], [op.deps[0].volume]).reshape(op.deps[0].shape);
69         auto packedGrads = dopt.core.ops.nnet.batchNormGrad(trimmedGrad, op.deps[0], op.deps[1]);
70         packedGrads = packedGrads.reshape([packedGrads.volume]);
71 
72         auto v0 = op.deps[0].volume;
73         auto v1 = op.deps[1].volume;
74         auto v2 = op.deps[2].volume;
75 
76         return [
77             packedGrads.slice([0], [v0]).reshape(op.deps[0].shape),
78             packedGrads.slice([v0], [v0 + v1]).reshape(op.deps[1].shape),
79             packedGrads.slice([v0 + v1], [v0 + v1 + v2]).reshape(op.deps[2].shape),
80             float32(op.deps[3].shape),
81             float32(op.deps[4].shape)
82         ];
83     }
84 }