Improve pruning module

paddle/parameter/ParameterUpdaterHook.cpp

@@ -14,11 +14,13 @@ limitations under the License. */
 
 #include "ParameterUpdaterHook.h"
 
+#include <algorithm>
 #include <atomic>
 #include <fstream>
 #include <mutex>
 #include <thread>
 #include <unordered_map>
+#include <vector>
 
 #include "paddle/math/Vector.h"
 #include "paddle/parameter/Parameter.h"
@@ -29,40 +31,21 @@ namespace paddle {
 
 /**
  * The static pruning hook
- *
- * Static means user load a mask map before training started. This map will
- * define which link/weight between neural is disabled.
+ * Static means user specific a sparsity_ratio before training start, and the
+ * network will prune the parameters based on the sparsity_ratio. More deatils
+ * can see https://arxiv.org/pdf/1506.02626.pdf.
  */
+
 class StaticPruningHook : public IParameterUpdaterHook {
 public:
-  /**
-   * The Mask Map Header.
-   * The map file started with this header.
-   *
-   * In Version 0, reset file will be:
-   *  contains header.size bit, each bit means such weight is enabled or not.
-   *    if bit is 1, then such weight is enabled.
-   *  at end, the file will round to byte, and the low bits of end byte will be
-   *  filled by zero.
-   *
-   */
-  struct StaticMaskHeader {
-    uint32_t version;
-    size_t size;
-  } __attribute__((__packed__));
-
-  explicit StaticPruningHook(const std::string& mask_filename) : initCount_(0) {
-    bool ok = this->loadMaskFile(mask_filename);
-    if (!ok) {
-      LOG(WARNING) << "Fail to load mask file " << mask_filename
-                   << " in current directory, searching in init_model_path";
-      std::string combineMaskFilename =
-          path::join(FLAGS_init_model_path, mask_filename);
-      CHECK(this->loadMaskFile(combineMaskFilename))
-          << "Cannot load " << mask_filename << " in ./" << mask_filename
-          << " and " << combineMaskFilename;
-    }
-    VLOG(3) << mask_filename << " mask size = " << this->mask_.size();
+  explicit StaticPruningHook(const ParameterUpdaterHookConfig& hookConfig)
+      : initCount_(0) {
+    sparsityRatio_ = hookConfig.sparsity_ratio();
+  }
+
+  static bool sortPairAscend(const std::pair<real, size_t>& pair1,
+                             const std::pair<real, size_t>& pair2) {
+    return pair1.first > pair2.first;
   }
 
   void update(Parameter* para) {
@@ -73,62 +56,53 @@ class StaticPruningHook : public IParameterUpdaterHook {
     }
   }
 
-  void init(Parameter* para) {
-    size_t initCount = this->initCount_.fetch_add(1);
-    CHECK_EQ(initCount, 0UL) << "Currently the StaticPruningHook must invoke "
-                                "in same ParamterUpdater";
-    VLOG(3) << "Initialize Parameter " << para;
-    SetDevice device(para->getDeviceId());
+  void generateMask(Parameter* para) {
+    VectorPtr vec = para->getBuf(PARAMETER_VALUE);
+    maskTemp_ = Vector::create(para->getSize(), false);
+    maskTemp_->zeroMem();
+    real* dataPtr = maskTemp_->getData();
+    size_t sparsityNum = para->getSize() * (1 - sparsityRatio_);
 
-    auto maskVec = Vector::create(this->mask_.size(), false);
-    {  // Initialize maskVec with float mask vector
-      real* dataPtr = maskVec->getData();
-      size_t i = 0;
-      for (bool m : mask_) {
-        dataPtr[i++] = m ? 1.0 : 0.0;
-      }
-    }
+    VectorPtr vecCpu = Vector::create(para->getSize(), false);
+    vecCpu->copyFrom(*vec);
+    std::vector<std::pair<real, size_t>> param;
+
+    for (size_t i = 0; i < para->getSize(); i++)
+      param.push_back(std::make_pair(fabs(vecCpu->getData()[i]), i));
+
+    std::partial_sort(param.begin(),
+                      param.begin() + sparsityNum,
+                      param.end(),
+                      sortPairAscend);
+    for (size_t i = 0; i < sparsityNum; i++) dataPtr[param[i].second] = 1.0;
 
     // Currently just use a mask vector for hack.
-    // @TODO(yuyang18): Implemented the mask operation in vector.
     if (para->useGpu()) {
-      maskVec_ = Vector::create(this->mask_.size(), para->useGpu());
-      maskVec_->copyFrom(*maskVec);
+      maskVec_ = Vector::create(para->getSize(), para->useGpu());
+      maskVec_->copyFrom(*maskTemp_);
     } else {
-      maskVec_ = maskVec;
+      maskVec_ = maskTemp_;
     }
+  }
+
+  void init(Parameter* para) {
+    generateMask(para);
+    size_t initCount = this->initCount_.fetch_add(1);
+    CHECK_EQ(initCount, 0UL) << "Currently the StaticPruningHook must invoke "
+                                "in same ParamterUpdater";
+    VLOG(3) << "Initialize Parameter " << para;
+    SetDevice device(para->getDeviceId());
 
     auto& vec = para->getBuf(PARAMETER_VALUE);
     vec->dotMul(*maskVec_);
   }
 
 private:
-  bool loadMaskFile(const std::string& mask_filename) {
-    std::ifstream fin;
-    fin.open(mask_filename);
-    if (fin.is_open()) {
-      StaticMaskHeader header;
-      fin.read(reinterpret_cast<char*>(&header), sizeof(StaticMaskHeader));
-      CHECK_EQ(header.version, 0UL);
-      mask_.resize(header.size);
-      uint8_t buf;
-      for (size_t i = 0; i < header.size; ++i, buf <<= 1) {
-        if (i % 8 == 0) {
-          fin.read(reinterpret_cast<char*>(&buf), sizeof(uint8_t));
-        }
-        mask_[i] = buf & 0x80;
-      }
-      fin.close();
-      return true;
-    } else {
-      return false;
-    }
-  }
-
   SameThreadChecker updateThreadChecker_;
   std::atomic<size_t> initCount_;
   VectorPtr maskVec_;
-  std::vector<bool> mask_;
+  VectorPtr maskTemp_;
+  real sparsityRatio_;
 };
 
 IParameterUpdaterHook::IParameterUpdaterHook() {}
@@ -166,10 +140,10 @@ static IParameterUpdaterHook* createImpl(
     const ParameterUpdaterHookConfig& config) {
   auto& type = config.type();
   if (type == "pruning") {
-    if (config.has_purning_mask_filename()) {
-      return new StaticPruningHook(config.purning_mask_filename());
-    }
+    return new StaticPruningHook(config);
   }
+
+  LOG(FATAL) << "Unknown Hook type:  " << type;
   return nullptr;
 }
 

proto/ParameterConfig.proto

@@ -25,8 +25,9 @@ enum ParameterInitStrategy {
 }
 
 message ParameterUpdaterHookConfig {
+  // hook type such as  'pruning'
   required string type = 1;
-  optional string purning_mask_filename = 2;
+  optional double sparsity_ratio = 2 [default = 0.8];
 }
 
 message ParameterConfig {

python/paddle/trainer/config_parser.py

@@ -3189,11 +3189,11 @@ def Layer(name, type, **xargs):
 @config_func
 def ParameterHook(type, **kwargs):
     if type == 'pruning':
-        mask_filename = kwargs.get('mask_filename', None)
-        assert mask_filename is not None
         hook = ParameterUpdaterHookConfig()
         hook.type = type
-        hook.purning_mask_filename = mask_filename
+        sparsity_ratio = kwargs.get('sparsity_ratio', None)
+        if sparsity_ratio is not None:
+            hook.sparsity_ratio = sparsity_ratio
         return hook
     else:
         return None
@@ -3300,13 +3300,13 @@ def Parameter(name,
 
     if update_hooks is not None:
         if hasattr(update_hooks, '__call__'):
-            update_hooks = update_hooks(para.name)
+            update_hooks = update_hooks()
 
         if isinstance(update_hooks, list):
             for hook in update_hooks:
                 para.update_hooks.extend([hook])
         else:
-            para.update_hooks.extend(update_hooks)
+            para.update_hooks.extend([update_hooks])
 
     g_parameter_map[name] = para
 

python/paddle/trainer_config_helpers/attrs.py

@@ -14,7 +14,8 @@
 
 from paddle.trainer.config_parser import *
 __all__ = [
-    'ParamAttr', 'ExtraAttr', 'ParameterAttribute', 'ExtraLayerAttribute'
+    'HookAttr', 'ParamAttr', 'ExtraAttr', 'ParameterAttribute',
+    'ExtraLayerAttribute'
 ]
 
 
@@ -55,6 +56,33 @@ def is_compatible_with(x, Type):
         return False
 
 
+class HookAttribute(object):
+    """
+    Hook Attribute object. The hook is an auxiliary operation that occurs 
+    during network propagation.
+    NOTE: IT IS A HIGH LEVEL USER INTERFACE.
+
+    :param  type: Hook type, eg: 'pruning'
+    :type type: string
+
+    :param sparsity_ratio: Must be specified if hook type is 'pruning'
+    :type sparsity_ratio: float or None
+
+    """
+
+    def __init__(self, type, sparsity_ratio=None):
+        self.type = type
+        self.sparsity_ratio = sparsity_ratio
+        if self.sparsity_ratio is not None:
+            assert is_compatible_with(
+                self.sparsity_ratio,
+                float), 'sparisity_ratio must be float type'
+            assert self.sparsity_ratio <= 1 and self.sparsity_ratio >= 0, 'sparisity must be a flaot between [0, 1] '
+
+    def __call__(self):
+        return ParameterHook(self.type, sparsity_ratio=self.sparsity_ratio)
+
+
 class ParameterAttribute(object):
     """
     Parameter Attributes object. To fine-tuning network training process, user
@@ -109,7 +137,8 @@ def __init__(self,
                  learning_rate=None,
                  momentum=None,
                  gradient_clipping_threshold=None,
-                 sparse_update=False):
+                 sparse_update=False,
+                 update_hooks=None):
         self.attr = {}
 
         if is_static:
@@ -162,6 +191,9 @@ def __init__(self,
             self.attr['gradient_clipping_threshold'] = \
                 gradient_clipping_threshold
 
+        if update_hooks:
+            self.attr['update_hooks'] = update_hooks
+
     def set_default_parameter_name(self, name):
         """
         Set default parameter name. If parameter not set, then will use default
@@ -237,5 +269,6 @@ def to_kwargs(attr):
             return attr.attr
 
 
+HookAttr = HookAttribute
 ParamAttr = ParameterAttribute
 ExtraAttr = ExtraLayerAttribute

python/paddle/v2/attr.py

@@ -17,10 +17,12 @@
 __all__ = [
     "Param",
     "Extra",
+    "Hook",
 ]
 
 Param = paddle.trainer_config_helpers.attrs.ParameterAttribute
 Extra = paddle.trainer_config_helpers.attrs.ExtraLayerAttribute
+Hook = paddle.trainer_config_helpers.attrs.HookAttribute
 
 for each in paddle.trainer_config_helpers.attrs.__all__:
     globals()[each] = getattr(paddle.trainer_config_helpers.attrs, each)