MKLDNNFcLayer.h 4.7 KB
Newer Older
T
tensor-tang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2017 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

17
#include "MKLDNNLayer.h"
T
tensor-tang 已提交
18 19 20
#include "mkldnn.hpp"

namespace paddle {
21 22 23
typedef mkldnn::inner_product_forward fc_fwd;
typedef mkldnn::inner_product_backward_weights fc_bwdWgt;
typedef mkldnn::inner_product_backward_data fc_bwdData;
T
tensor-tang 已提交
24 25

/**
26
 * @brief A subclass of MKLDNNLayer fc layer.
T
tensor-tang 已提交
27 28 29
 *
 * The config file api is mkldnn_fc
 */
30
class MKLDNNFcLayer : public MKLDNNLayer {
T
tensor-tang 已提交
31
protected:
T
tensor-tang 已提交
32 33 34
  // input layer size, can not be change after init
  size_t iLayerSize_;  // == ic * ih * iw

T
tensor-tang 已提交
35
  // if has already init the weight
T
tensor-tang 已提交
36
  bool hasInitedWgt_;
T
tensor-tang 已提交
37

38 39 40
  // save forward primitive_desc, which can be used backward
  std::shared_ptr<fc_fwd::primitive_desc> fwdPD_;

T
tensor-tang 已提交
41 42 43 44
  // fc weight and bias
  std::unique_ptr<Weight> weight_;
  std::unique_ptr<Weight> biases_;

T
tensor-tang 已提交
45
public:
46
  explicit MKLDNNFcLayer(const LayerConfig& config)
47
      : MKLDNNLayer(config), hasInitedWgt_(false) {}
T
tensor-tang 已提交
48

49
  ~MKLDNNFcLayer() {}
T
tensor-tang 已提交
50 51 52 53

  bool init(const LayerMap& layerMap,
            const ParameterMap& parameterMap) override;

54 55 56 57 58 59 60 61 62 63 64 65 66 67
  void reshape(
      int& bs, int& ic, int& ih, int& iw, int oc, int& oh, int& ow) override;

  void resetFwd(std::vector<mkldnn::primitive>& pipeline,
                MKLDNNMatrixPtr& in,
                MKLDNNMatrixPtr& wgt,
                MKLDNNMatrixPtr& bias,
                MKLDNNMatrixPtr& out) override;

  void resetBwd(std::vector<mkldnn::primitive>& pipeline,
                MKLDNNMatrixPtr& in,
                MKLDNNMatrixPtr& wgt,
                MKLDNNMatrixPtr& bias,
                MKLDNNMatrixPtr& out) override;
T
tensor-tang 已提交
68

69 70
  void updateInputData() override;

71 72 73 74 75
  void updateWeights(const UpdateCallback& callback) override;

  void convertWeightsFromPaddle() override;

  void convertWeightsToPaddle() override;
76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128

protected:
  /**
   * Forward functions: reset buffers(input, output, weight and bias),
   *                    reset primitive descriptor,
   *                    reset pipeline.
   */
  void resetFwdBuffers(MKLDNNMatrixPtr& in,
                       MKLDNNMatrixPtr& wgt,
                       MKLDNNMatrixPtr& bias,
                       MKLDNNMatrixPtr& out);
  void resetInValue(MKLDNNMatrixPtr& in);
  void resetWgtBiasValue(MKLDNNMatrixPtr& wgt, MKLDNNMatrixPtr& bias);
  void resetOutValue(MKLDNNMatrixPtr& out);
  void resetFwdPD(std::shared_ptr<fc_fwd::primitive_desc>& pd,
                  MKLDNNMatrixPtr in,
                  MKLDNNMatrixPtr wgt,
                  MKLDNNMatrixPtr bias,
                  MKLDNNMatrixPtr out);
  void resetFwdPipeline(std::vector<mkldnn::primitive>& pipeline,
                        std::shared_ptr<fc_fwd::primitive_desc>& pd,
                        MKLDNNMatrixPtr& in,
                        MKLDNNMatrixPtr& wgt,
                        MKLDNNMatrixPtr& bias,
                        MKLDNNMatrixPtr& out);

  /**
   * Backward functions: reset buffers(input, output, weight and bias),
   *                     reset primitive descriptor for backward weight,
   *                     reset primitive descriptor for backward data,
   *                     reset pipeline.
   */
  void resetBwdBuffers(MKLDNNMatrixPtr& in,
                       MKLDNNMatrixPtr& wgt,
                       MKLDNNMatrixPtr& bias,
                       MKLDNNMatrixPtr& out);
  void resetOutGrad(MKLDNNMatrixPtr& out);
  void resetWgtBiasGrad(MKLDNNMatrixPtr& wgt, MKLDNNMatrixPtr& bias);
  void resetInGrad(MKLDNNMatrixPtr& in);
  void resetBwdWgtPD(std::shared_ptr<fc_bwdWgt::primitive_desc>& pd,
                     MKLDNNMatrixPtr& wgt,
                     MKLDNNMatrixPtr& bias,
                     MKLDNNMatrixPtr& out);
  void resetBwdDataPD(std::shared_ptr<fc_bwdData::primitive_desc>& pd,
                      MKLDNNMatrixPtr& in,
                      MKLDNNMatrixPtr& out);
  void resetBwdPipeline(std::vector<mkldnn::primitive>& pipeline,
                        std::shared_ptr<fc_bwdWgt::primitive_desc>& bwdWgtPD,
                        std::shared_ptr<fc_bwdData::primitive_desc>& bwdDataPD,
                        MKLDNNMatrixPtr& in,
                        MKLDNNMatrixPtr& wgt,
                        MKLDNNMatrixPtr& bias,
                        MKLDNNMatrixPtr& out);
T
tensor-tang 已提交
129 130 131
};

}  // namespace paddle