Line data    Source code

       1             : #include "DnnlSoftmaxLayer.h"
       2             : 
       3             : namespace elsa::ml
       4             : {
       5             :     namespace detail
       6             :     {
       7             :         template <typename data_t>
       8           0 :         DnnlSoftmaxLayer<data_t>::DnnlSoftmaxLayer(const VolumeDescriptor& inputDescriptor,
       9             :                                                    const VolumeDescriptor& outputDescriptor)
      10           0 :             : DnnlLayer<data_t>(inputDescriptor, outputDescriptor, "DnnlSoftmaxLayer")
      11             :         {
      12           0 :             _softmaxAxis = (inputDescriptor.getNumberOfDimensions() == 1) ? 0 : 1;
      13           0 :         }
      14             : 
      15             :         template <typename data_t>
      16           0 :         void DnnlSoftmaxLayer<data_t>::compileForwardStream()
      17             :         {
      18           0 :             BaseType::compileForwardStream();
      19             : 
      20           0 :             auto desc = dnnl::softmax_forward::desc(dnnl::prop_kind::forward_training,
      21           0 :                                                     _input.front().descriptor, _softmaxAxis);
      22             : 
      23           0 :             _forwardPrimitiveDescriptor = dnnl::softmax_forward::primitive_desc(desc, *_engine);
      24             : 
      25             :             // Set forward primitive
      26           0 :             ELSA_ML_ADD_DNNL_PRIMITIVE(_forwardStream,
      27             :                                        dnnl::softmax_forward(_forwardPrimitiveDescriptor));
      28             : 
      29           0 :             _output.effectiveMemory =
      30           0 :                 std::make_shared<dnnl::memory>(_forwardPrimitiveDescriptor.dst_desc(), *_engine);
      31             : 
      32           0 :             BaseType::validateDnnlMemory(_input.front().effectiveMemory, _output.effectiveMemory);
      33             : 
      34           0 :             _forwardStream.arguments.push_back({{DNNL_ARG_SRC, *_input.front().effectiveMemory},
      35           0 :                                                 {DNNL_ARG_DST, *_output.effectiveMemory}});
      36           0 :         }
      37             : 
      38             :         template <typename data_t>
      39           0 :         void DnnlSoftmaxLayer<data_t>::compileBackwardStream()
      40             :         {
      41           0 :             BaseType::compileBackwardStream();
      42             : 
      43           0 :             auto desc = dnnl::softmax_backward::desc(
      44           0 :                 /* Output gradient descriptor */ _outputGradient.front().descriptor,
      45           0 :                 /* Input descriptor */ _input.front().descriptor,
      46             :                 /* Softmax axis */ _softmaxAxis);
      47             : 
      48           0 :             _backwardPrimitiveDescriptor =
      49           0 :                 dnnl::softmax_backward::primitive_desc(desc, *_engine, _forwardPrimitiveDescriptor);
      50             : 
      51             :             // Reorder if necessary
      52           0 :             this->reorderMemory(_backwardPrimitiveDescriptor.diff_dst_desc(),
      53           0 :                                 _outputGradient.front(), _backwardStream);
      54             : 
      55           0 :             _inputGradient.front().effectiveMemory = std::make_shared<dnnl::memory>(
      56           0 :                 _backwardPrimitiveDescriptor.diff_src_desc(), *_engine);
      57             : 
      58           0 :             BaseType::validateDnnlMemory(_input.front().effectiveMemory, _output.effectiveMemory,
      59           0 :                                          _outputGradient.front().effectiveMemory,
      60           0 :                                          _inputGradient.front().effectiveMemory);
      61             : 
      62           0 :             ELSA_ML_ADD_DNNL_PRIMITIVE(_backwardStream,
      63             :                                        dnnl::softmax_backward(_backwardPrimitiveDescriptor));
      64           0 :             _backwardStream.arguments.push_back(
      65           0 :                 {{DNNL_ARG_DST, *_output.effectiveMemory},
      66           0 :                  {DNNL_ARG_DIFF_DST, *_outputGradient.front().effectiveMemory},
      67           0 :                  {DNNL_ARG_DIFF_SRC, *_inputGradient.front().effectiveMemory}});
      68           0 :         }
      69             : 
      70             :         template class DnnlSoftmaxLayer<float>;
      71             :     } // namespace detail
      72             : } // namespace elsa::ml