ParameterUpdaterHook.cpp 4.6 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Z
zhangjinchao01 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

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. */

#include "ParameterUpdaterHook.h"

Z
Zhaolong Xing 已提交
17
#include <algorithm>
Y
Yu Yang 已提交
18
#include <atomic>
Z
zhangjinchao01 已提交
19 20 21
#include <fstream>
#include <mutex>
#include <thread>
Y
Yu Yang 已提交
22
#include <unordered_map>
23
#include <vector>
Z
zhangjinchao01 已提交
24 25 26 27

#include "paddle/math/Vector.h"
#include "paddle/parameter/Parameter.h"
#include "paddle/utils/Flags.h"
Y
Yu Yang 已提交
28
#include "paddle/utils/Util.h"
Z
zhangjinchao01 已提交
29 30 31 32 33

namespace paddle {

/**
 * The static pruning hook
34
 * Static means user specify a sparsity_ratio before training started, and the
Z
zlx 已提交
35
 * network will prune the parameters based on the sparsity_ratio. More details
36
 * can be found https://arxiv.org/pdf/1506.02626.pdf.
Z
zhangjinchao01 已提交
37 38
 */

39
class StaticPruningHook : public IParameterUpdaterHook {
W
Wu Yi 已提交
40
 public:
Z
zlx 已提交
41
  explicit StaticPruningHook(const ParameterUpdaterHookConfig &hookConfig)
X
xzl 已提交
42 43 44 45
      : initCount_(0) {
    sparsityRatio_ = hookConfig.sparsity_ratio();
  }

Z
zlx 已提交
46 47
  static bool sortPairAscend(const std::pair<real, size_t> &pair1,
                             const std::pair<real, size_t> &pair2) {
X
xzl 已提交
48 49 50
    return pair1.first > pair2.first;
  }

Z
zlx 已提交
51
  void update(Parameter *para) {
X
xzl 已提交
52
    updateThreadChecker_.check();
Z
zlx 已提交
53
    auto &vec = para->getBuf(PARAMETER_GRADIENT);
X
xzl 已提交
54 55 56 57 58
    if (vec) {
      vec->dotMul(*maskVec_);
    }
  }

Z
zlx 已提交
59
  void generateMask(Parameter *para) {
60 61
    VectorPtr maskTemp = Vector::create(para->getSize(), false);
    maskTemp->zeroMem();
Z
zlx 已提交
62
    real *maskTempData = maskTemp->getData();
Z
zlx 已提交
63
    size_t nonZeroNum = para->getSize() * (1 - sparsityRatio_);
X
xzl 已提交
64

65 66 67 68
    VectorPtr paraVec = para->getBuf(PARAMETER_VALUE);
    VectorPtr paraCpuCopy = Vector::create(para->getSize(), false);

    paraCpuCopy->copyFrom(*paraVec);
69
    std::vector<std::pair<real, size_t>> param;
X
xzl 已提交
70 71

    for (size_t i = 0; i < para->getSize(); i++)
72
      param.push_back(std::make_pair(fabs(paraCpuCopy->getData()[i]), i));
X
xzl 已提交
73

Z
zlx 已提交
74 75 76
    std::partial_sort(
        param.begin(), param.begin() + nonZeroNum, param.end(), sortPairAscend);
    for (size_t i = 0; i < nonZeroNum; i++) maskTempData[param[i].second] = 1.0;
X
xzl 已提交
77 78 79 80

    // Currently just use a mask vector for hack.
    if (para->useGpu()) {
      maskVec_ = Vector::create(para->getSize(), para->useGpu());
81
      maskVec_->copyFrom(*maskTemp);
X
xzl 已提交
82
    } else {
83
      maskVec_ = maskTemp;
X
xzl 已提交
84
    }
X
xzl 已提交
85 86
  }

Z
zlx 已提交
87
  void init(Parameter *para) {
X
xzl 已提交
88 89
    generateMask(para);
    size_t initCount = this->initCount_.fetch_add(1);
90
    CHECK_EQ(initCount, 0UL) << "Currently the StaticPruningHook must invoke "
X
xzl 已提交
91 92 93 94
                                "in same ParamterUpdater";
    VLOG(3) << "Initialize Parameter " << para;
    SetDevice device(para->getDeviceId());

Z
zlx 已提交
95
    auto &paraVec = para->getBuf(PARAMETER_VALUE);
96
    paraVec->dotMul(*maskVec_);
X
xzl 已提交
97 98
  }

W
Wu Yi 已提交
99
 private:
X
xzl 已提交
100 101 102 103 104 105
  SameThreadChecker updateThreadChecker_;
  std::atomic<size_t> initCount_;
  VectorPtr maskVec_;
  real sparsityRatio_;
};

Z
zhangjinchao01 已提交
106 107 108 109 110 111 112 113 114 115 116 117 118
IParameterUpdaterHook::IParameterUpdaterHook() {}

IParameterUpdaterHook::~IParameterUpdaterHook() {}

/**
 * A Hasher used by g_hooks.
 *
 * Use the independent hasher intendedly. There is a hasher in PServer for hash
 * ParameterBlock. But not to use same hasher to reduce dependency.
 *
 * May be extracted to Util.h to unify the hasher.
 */
class StringIntPairHasher {
W
Wu Yi 已提交
119
 public:
Z
zlx 已提交
120
  size_t operator()(const std::pair<std::string, int> &k) const {
Z
zhangjinchao01 已提交
121 122 123
    return intHasher_(strHasher_(k.first) + k.second);
  }

W
Wu Yi 已提交
124
 private:
Z
zhangjinchao01 已提交
125 126 127 128
  std::hash<std::string> strHasher_;
  std::hash<int> intHasher_;
};

Z
zlx 已提交
129 130 131 132
static WeakKVCache<std::pair<std::string, int>,
                   IParameterUpdaterHook,
                   StringIntPairHasher>
    g_hookCache_;
Z
zhangjinchao01 已提交
133 134 135 136

/**
 * ParameterUpdaterHook actually factory method.
 */
Z
zlx 已提交
137 138
static IParameterUpdaterHook *createImpl(
    const ParameterUpdaterHookConfig &config) {
Z
zlx 已提交
139
  auto &type = config.type();
140
  if (type == "pruning") {
X
xzl 已提交
141
    return new StaticPruningHook(config);
Z
zhangjinchao01 已提交
142
  }
X
xzl 已提交
143 144

  LOG(FATAL) << "Unknown Hook type:  " << type;
Z
zhangjinchao01 已提交
145 146 147
  return nullptr;
}

Z
zlx 已提交
148 149
std::shared_ptr<IParameterUpdaterHook> IParameterUpdaterHook::create(
    const ParameterConfig &paramConfig, int idx) {
Z
zhangjinchao01 已提交
150 151 152 153 154
  std::pair<std::string, int> key = {paramConfig.name(), idx};
  return g_hookCache_.get(
      key, [&] { return createImpl(paramConfig.update_hooks(idx)); });
}

Z
zlx 已提交
155
}  // namespace paddle