Line data Source code
1 : #include "DnnlDenseLayer.h"
2 :
3 : namespace elsa::ml
4 : {
5 : namespace detail
6 : {
7 : template <typename data_t>
8 0 : DnnlDenseLayer<data_t>::DnnlDenseLayer(const VolumeDescriptor& inputDescriptor,
9 : const VolumeDescriptor& outputDescriptor,
10 : const VolumeDescriptor& weightsDescriptor,
11 : Initializer initializer)
12 : : DnnlTrainableLayer<data_t>(inputDescriptor, outputDescriptor, weightsDescriptor,
13 0 : initializer)
14 : {
15 0 : }
16 :
17 : template <typename data_t>
18 0 : void DnnlDenseLayer<data_t>::compileForwardStream()
19 : {
20 0 : BaseType::compileForwardStream();
21 :
22 0 : auto desc = dnnl::inner_product_forward::desc(
23 : /* Propagation kind */ dnnl::prop_kind::forward_training,
24 0 : /* Source descriptor*/ _input.front().descriptor,
25 0 : /* Weights memory descriptor*/ _weights.descriptor,
26 0 : /* Bias memory descriptor*/ _bias.descriptor,
27 0 : /* Destination memorydescriptor*/ _output.descriptor);
28 :
29 : // Create inner-product forward primitive
30 0 : _forwardPrimitiveDescriptor =
31 0 : dnnl::inner_product_forward::primitive_desc(desc, *_engine);
32 :
33 : // Do we need to reorder? This is the case of the memory description chosen by the
34 : // primitive differs from the memory description we provided based on the layer's input
35 : // descriptor / weights descriptor.
36 0 : this->reorderMemory(_forwardPrimitiveDescriptor.src_desc(), _input.front(),
37 0 : _forwardStream);
38 0 : this->reorderMemory(_forwardPrimitiveDescriptor.weights_desc(), _weights,
39 0 : _forwardStream);
40 :
41 : // Allocate output-demmory
42 0 : _output.effectiveMemory =
43 0 : std::make_shared<dnnl::memory>(_forwardPrimitiveDescriptor.dst_desc(), *_engine);
44 :
45 : // Add inner-product primitive to forward-stream
46 0 : ELSA_ML_ADD_DNNL_PRIMITIVE(_forwardStream,
47 : dnnl::inner_product_forward(_forwardPrimitiveDescriptor));
48 :
49 : // Validate memory
50 0 : BaseType::validateDnnlMemory(_input.front().effectiveMemory, _weights.effectiveMemory,
51 0 : _bias.effectiveMemory, _output.effectiveMemory);
52 :
53 0 : _forwardStream.arguments.push_back({{DNNL_ARG_SRC, *_input.front().effectiveMemory},
54 0 : {DNNL_ARG_WEIGHTS, *_weights.effectiveMemory},
55 0 : {DNNL_ARG_BIAS, *_bias.effectiveMemory},
56 0 : {DNNL_ARG_DST, *_output.effectiveMemory}});
57 0 : }
58 :
59 : template <typename data_t>
60 0 : void DnnlDenseLayer<data_t>::compileBackwardStream()
61 : {
62 0 : BaseType::compileBackwardStream();
63 0 : compileBackwardWeightsStream();
64 0 : compileBackwardDataStream();
65 0 : }
66 :
67 : template <typename data_t>
68 0 : void DnnlDenseLayer<data_t>::compileBackwardDataStream()
69 : {
70 0 : auto desc = dnnl::inner_product_backward_data::desc(
71 0 : /* Input gradient descriptor (output) */ _inputGradient.front().descriptor,
72 0 : /* Weights descriptor */ _weights.descriptor,
73 0 : /* Output gradient descriptor */ _outputGradient.front().descriptor);
74 :
75 0 : _backwardDataPrimitiveDescriptor = dnnl::inner_product_backward_data::primitive_desc(
76 0 : desc, *_engine, _forwardPrimitiveDescriptor);
77 :
78 : // Reorder output gradient of necessary
79 0 : this->reorderMemory(_backwardDataPrimitiveDescriptor.diff_dst_desc(),
80 0 : _outputGradient.front(), _backwardStream);
81 :
82 : // Reorder weights if necessary
83 0 : this->reorderMemory(_backwardDataPrimitiveDescriptor.weights_desc(), _weights,
84 0 : _backwardStream);
85 :
86 : // Set input gradient memory
87 0 : _inputGradient.front().effectiveMemory = std::make_shared<dnnl::memory>(
88 0 : _backwardDataPrimitiveDescriptor.diff_src_desc(), *_engine);
89 :
90 : // Push backward data primitive
91 0 : ELSA_ML_ADD_DNNL_PRIMITIVE(_backwardStream, dnnl::inner_product_backward_data(
92 : _backwardDataPrimitiveDescriptor));
93 :
94 0 : BaseType::validateDnnlMemory(_inputGradient.front().effectiveMemory,
95 0 : _weights.effectiveMemory,
96 0 : _outputGradient.front().effectiveMemory);
97 :
98 0 : _backwardStream.arguments.push_back(
99 0 : {/* Input gradient */ {DNNL_ARG_DIFF_SRC, *_inputGradient.front().effectiveMemory},
100 0 : /* Weights */ {DNNL_ARG_WEIGHTS, *_weights.effectiveMemory},
101 : /* Output gradient */
102 0 : {DNNL_ARG_DIFF_DST, *_outputGradient.front().effectiveMemory}});
103 0 : }
104 :
105 : template <typename data_t>
106 0 : void DnnlDenseLayer<data_t>::compileBackwardWeightsStream()
107 : {
108 : // Backward descriptor for weights backprop
109 0 : auto desc = dnnl::inner_product_backward_weights::desc(
110 0 : /* Input descriptor */ _input.front().descriptor,
111 0 : /* Weights gradient descriptor */ _weightsGradient.descriptor,
112 0 : /* Bias gradient descriptor */ _biasGradient.descriptor,
113 0 : /* Output gradient descriptor */ _outputGradient.front().descriptor);
114 :
115 0 : _backwardWeightsPrimitiveDescriptor =
116 0 : dnnl::inner_product_backward_weights::primitive_desc(desc, *_engine,
117 0 : _forwardPrimitiveDescriptor);
118 :
119 : // Do we need reorder for gradient src memory?
120 0 : this->reorderMemory(_backwardWeightsPrimitiveDescriptor.src_desc(), _input.front(),
121 0 : _backwardStream);
122 :
123 0 : this->reorderMemory(_backwardWeightsPrimitiveDescriptor.diff_dst_desc(),
124 0 : _outputGradient.front(), _backwardStream);
125 :
126 0 : BaseType::validateDnnlMemory(_input.front().effectiveMemory,
127 0 : _biasGradient.effectiveMemory,
128 0 : _outputGradient.front().effectiveMemory);
129 :
130 0 : ELSA_ML_ADD_DNNL_PRIMITIVE(_backwardStream, dnnl::inner_product_backward_weights(
131 : _backwardWeightsPrimitiveDescriptor));
132 0 : _backwardStream.arguments.push_back(
133 0 : {{DNNL_ARG_SRC, *_input.front().effectiveMemory},
134 0 : {DNNL_ARG_DIFF_BIAS, *_biasGradient.effectiveMemory},
135 0 : {DNNL_ARG_DIFF_DST, *_outputGradient.front().effectiveMemory}});
136 :
137 0 : _weightsGradient.effectiveMemory = _weightsGradient.describedMemory;
138 0 : if (_weightsGradient.describedMemory->get_desc()
139 0 : != _backwardWeightsPrimitiveDescriptor.diff_weights_desc()) {
140 0 : _weightsGradient.wasReordered = true;
141 0 : _weightsGradient.describedMemory = std::make_shared<dnnl::memory>(
142 0 : _backwardWeightsPrimitiveDescriptor.diff_weights_desc(), *_engine);
143 0 : _backwardStream.arguments.back().insert(
144 0 : {DNNL_ARG_DIFF_WEIGHTS, *_weightsGradient.describedMemory});
145 0 : ELSA_ML_ADD_DNNL_PRIMITIVE(_backwardStream,
146 : dnnl::reorder(*_weightsGradient.describedMemory,
147 : *_weightsGradient.effectiveMemory));
148 0 : _backwardStream.arguments.push_back(
149 0 : {{DNNL_ARG_FROM, *_weightsGradient.describedMemory},
150 0 : {DNNL_ARG_TO, *_weightsGradient.effectiveMemory}});
151 : } else {
152 0 : _backwardStream.arguments.back().insert(
153 0 : {DNNL_ARG_DIFF_WEIGHTS, *_weightsGradient.effectiveMemory});
154 : }
155 0 : }
156 :
157 : template class DnnlDenseLayer<float>;
158 : } // namespace detail
159 : } // namespace elsa::ml
|