MKLDNNFcLayer.h 3.8 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
  void reshape(
55
      int& bs, int& ic, int& ih, int& iw, int& oc, int& oh, int& ow) override;
56 57

  void resetFwd(std::vector<mkldnn::primitive>& pipeline,
58
                std::vector<MKLDNNMatrixPtr>& inputs,
59 60 61
                MKLDNNMatrixPtr& out) override;

  void resetBwd(std::vector<mkldnn::primitive>& pipeline,
62
                std::vector<MKLDNNMatrixPtr>& inputs,
63
                MKLDNNMatrixPtr& out) override;
T
tensor-tang 已提交
64

65 66 67 68 69
  void updateWeights(const UpdateCallback& callback) override;

  void convertWeightsFromPaddle() override;

  void convertWeightsToPaddle() override;
70 71 72 73 74 75 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

protected:
  void resetFwdBuffers(MKLDNNMatrixPtr& in,
                       MKLDNNMatrixPtr& wgt,
                       MKLDNNMatrixPtr& bias,
                       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);
  void resetBwdBuffers(MKLDNNMatrixPtr& in,
                       MKLDNNMatrixPtr& wgt,
                       MKLDNNMatrixPtr& bias,
                       MKLDNNMatrixPtr& out);
  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 已提交
105 106 107
};

}  // namespace paddle