[TensorV2] Create a mask of dropout, filter, and zoneout

This PR adds new functionalities for getting masks of the following techniques: dropout, filter, and zoneout.
These functions enable working with masks, making it easier to perform such techniques in regularization.

**Changes proposed in this PR:**
- Added dropout_mask(), which calculates the dropout mask by multiplying tensor elements by 1.0 / (1.0 - dropout rate), in place.
- Added filter_mask(), which takes an input tensor and applies a filter mask based on the given mask length and invert option.
- Added zoneout_mask(), which generates two zoneout masks, one for in-place operation and another for opposite operation, based on the specified zoneout rate.

**Self-evaluation:**
1. Build test: [X]Passed [ ]Failed [ ]Skipped
2. Run test:   [X]Passed [ ]Failed [ ]Skipped

Signed-off-by: Donghyeon Jeong <[email protected]>

nntrainer/tensor/float_tensor.cpp

@@ -634,6 +634,53 @@ TensorV2 &FloatTensor::dot(TensorV2 const &input, TensorV2 &output, bool trans,
   return output;
 }
 
+void FloatTensor::dropout_mask(float dropout) {
+  float scale = 1.0 / (1 - dropout);
+  float *data_ = (float *)getData();
+  for (unsigned int i = 0; i < size(); ++i) {
+    if (data_[i] >= dropout)
+      data_[i] = scale;
+    else
+      data_[i] = 0.0;
+  }
+}
+
+void FloatTensor::filter_mask(const TensorV2 &mask_len, bool reverse) {
+  float fill_mask_val = 0.0;
+  float en_mask_val = 1.0 - fill_mask_val;
+
+  if (reverse) {
+    fill_mask_val = 1.0;
+    en_mask_val = 1.0 - fill_mask_val;
+  }
+
+  setValue(fill_mask_val);
+
+  NNTR_THROW_IF(mask_len.batch() != batch(), std::invalid_argument)
+    << "Number of filter masks mismatched";
+
+  for (unsigned int b = 0; b < batch(); b++) {
+    float *addr = (float *)getAddress(getIndex(b, 0, 0, 0));
+    const uint *mask_len_val = mask_len.getAddress<uint>(b, 0, 0, 0);
+    std::fill(addr, addr + (*mask_len_val), en_mask_val);
+  }
+}
+
+void FloatTensor::zoneout_mask(TensorV2 &opposite, float zoneout) {
+  opposite.setRandBernoulli(zoneout);
+
+  float *data = (float *)getData();
+  float *opposite_data = opposite.getData<float>();
+
+  for (unsigned int i = 0; i < size(); ++i) {
+    if (opposite_data[i] > epsilon) {
+      data[i] = 0.0f;
+    } else {
+      data[i] = 1.0f;
+    }
+  }
+}
+
 void FloatTensor::print(std::ostream &out) const {
   printInstance(out, this);
   const float *data = (float *)getData();

nntrainer/tensor/float_tensor.h

@@ -275,6 +275,21 @@ class FloatTensor : public TensorBase {
   TensorV2 &dot(TensorV2 const &input, TensorV2 &output, bool trans,
                 bool trans_in, float beta) const override;
 
+  /**
+   * @copydoc TensorV2::dropout_mask(float dropout)
+   */
+  void dropout_mask(float dropout) override;
+
+  /**
+   * @copydoc TensorV2::filter_mask(const TensorV2 &mask_len, bool reverse)
+   */
+  void filter_mask(const TensorV2 &mask_len, bool reverse) override;
+
+  /**
+   * @copydoc TensorV2::zoneout_mask(TensorV2 &opposite, float zoneout)
+   */
+  void zoneout_mask(TensorV2 &opposite, float zoneout) override;
+
   /**
    * @copydoc TensorV2::copy(const TensorV2 &from)
    */

nntrainer/tensor/half_tensor.cpp

@@ -607,6 +607,54 @@ TensorV2 &HalfTensor::dot(TensorV2 const &input, TensorV2 &output, bool trans,
   return output;
 }
 
+void HalfTensor::dropout_mask(float dropout) {
+  _FP16 scale = static_cast<_FP16>(1.0 / (1 - dropout));
+  _FP16 *data_ = (_FP16 *)getData();
+  for (unsigned int i = 0; i < size(); ++i) {
+    if (data_[i] >= dropout)
+      data_[i] = scale;
+    else
+      data_[i] = 0;
+  }
+}
+
+void HalfTensor::filter_mask(const TensorV2 &mask_len, bool reverse) {
+  float fill_mask_val = 0.0;
+  float en_mask_val = 1.0 - fill_mask_val;
+
+  if (reverse) {
+    fill_mask_val = 1.0;
+    en_mask_val = 1.0 - fill_mask_val;
+  }
+
+  setValue(fill_mask_val);
+
+  NNTR_THROW_IF(mask_len.batch() != batch(), std::invalid_argument)
+    << "Number of filter masks mismatched";
+
+  for (unsigned int b = 0; b < batch(); b++) {
+    _FP16 *addr = (_FP16 *)getAddress(getIndex(b, 0, 0, 0));
+    const uint *mask_len_val = mask_len.getAddress<uint>(b, 0, 0, 0);
+    std::fill(addr, addr + (*mask_len_val), (_FP16)en_mask_val);
+  }
+}
+
+void HalfTensor::zoneout_mask(TensorV2 &opposite, float zoneout) {
+  _FP16 zoneout_fp16 = (_FP16)zoneout;
+  opposite.setRandBernoulli(zoneout_fp16);
+
+  _FP16 *data = (_FP16 *)getData();
+  _FP16 *opposite_data = opposite.getData<_FP16>();
+
+  for (unsigned int i = 0; i < size(); ++i) {
+    if (opposite_data[i] > epsilon) {
+      data[i] = (_FP16)0.0;
+    } else {
+      data[i] = (_FP16)1.0;
+    }
+  }
+}
+
 void HalfTensor::print(std::ostream &out) const {
   printInstance(out, this);
   const _FP16 *data = (_FP16 *)getData();

nntrainer/tensor/half_tensor.h

@@ -274,6 +274,21 @@ class HalfTensor : public TensorBase {
   TensorV2 &dot(TensorV2 const &input, TensorV2 &output, bool trans,
                 bool trans_in, float beta) const override;
 
+  /**
+   * @copydoc TensorV2::dropout_mask(float dropout)
+   */
+  void dropout_mask(float dropout) override;
+
+  /**
+   * @copydoc TensorV2::filter_mask(const TensorV2 &mask_len, bool reverse)
+   */
+  void filter_mask(const TensorV2 &mask_len, bool reverse) override;
+
+  /**
+   * @copydoc TensorV2::zoneout_mask(TensorV2 &opposite, float zoneout)
+   */
+  void zoneout_mask(TensorV2 &opposite, float zoneout) override;
+
   /**
    * @copydoc TensorV2::copy(const TensorV2 &from)
    */

nntrainer/tensor/tensor_base.h

@@ -270,6 +270,21 @@ class TensorBase {
   virtual TensorV2 &dot(TensorV2 const &input, TensorV2 &output, bool trans,
                         bool trans_in, float beta) const = 0;
 
+  /**
+   * @copydoc TensorV2::dropout_mask(float dropout)
+   */
+  virtual void dropout_mask(float dropout) = 0;
+
+  /**
+   * @copydoc TensorV2::filter_mask(const TensorV2 &mask_len, bool reverse)
+   */
+  virtual void filter_mask(const TensorV2 &mask_len, bool reverse) = 0;
+
+  /**
+   * @copydoc TensorV2::zoneout_mask(TensorV2 &opposite, float zoneout)
+   */
+  virtual void zoneout_mask(TensorV2 &opposite, float zoneout) = 0;
+
   /**
    * @copydoc TensorV2::print(std::ostream &out)
    */
@@ -557,8 +572,7 @@ class SrcSharedTensorBase {
    * @brief   Constructor for the class
    */
   SrcSharedTensorBase(const TensorBase *tensor, size_t offset) :
-    src(tensor),
-    off(offset) {}
+    src(tensor), off(offset) {}
 
   /**
    * @brief   Get the allocated src tensor

nntrainer/tensor/tensor_v2.cpp

@@ -507,6 +507,50 @@ TensorV2 &TensorV2::dot_batched_deriv_wrt_2(TensorV2 &m_deriv,
   }
 }
 
+TensorV2 TensorV2::dropout_mask(float dropout) const {
+  TensorV2 output(getDim());
+  output.dropout_mask(dropout);
+  return output;
+}
+
+void TensorV2::dropout_mask(float dropout) {
+  /// @todo add unittest
+  NNTR_THROW_IF(dropout < 0 || dropout > 1, std::invalid_argument)
+    << "[Tensor::dropout_mask] Dropout rate should be between 0 and 1";
+
+  // if the rate is zero, no change is needed
+  if (std::fpclassify(dropout) == FP_ZERO)
+    return;
+
+  setRandUniform(0.0, 1.0);
+  itensor->dropout_mask(dropout);
+}
+
+void TensorV2::filter_mask(const TensorV2 &mask_len, bool reverse) {
+  /// @todo add unittest
+  itensor->filter_mask(mask_len, reverse);
+}
+
+TensorV2 TensorV2::zoneout_mask(float zoneout) {
+  TensorV2 output(getDim());
+  zoneout_mask(output, zoneout);
+  return output;
+}
+
+void TensorV2::zoneout_mask(TensorV2 &opposite, float zoneout) {
+  NNTR_THROW_IF(getDim() != opposite.getDim(), std::invalid_argument)
+    << "[Tensor::zoneout_mask] opposite dimension does not match";
+
+  NNTR_THROW_IF(zoneout < 0 || zoneout > 1, std::invalid_argument)
+    << "[Tensor::zoneout_mask] Zoneout rate should be between 0 and 1";
+
+  // if the rate is zero, no change is needed
+  if (std::fpclassify(zoneout) == FP_ZERO)
+    return;
+
+  itensor->zoneout_mask(opposite, zoneout);
+}
+
 void TensorV2::print(std::ostream &out) const { itensor->print(out); }
 
 void TensorV2::putData() const { itensor->putData(); }
@@ -596,8 +640,8 @@ const bool TensorV2::getContiguous() const noexcept {
   return itensor->getContiguous();
 }
 
-const std::array<size_t, TensorDim::MAXDIM> TensorV2::getStrides() const
-  noexcept {
+const std::array<size_t, TensorDim::MAXDIM>
+TensorV2::getStrides() const noexcept {
   return itensor->getStrides();
 }
 

nntrainer/tensor/tensor_v2.h

@@ -874,6 +874,46 @@ class TensorV2 {
                                     bool trans = false, bool trans_in = false,
                                     float beta = 0.0f) const;
 
+  /**
+   * @brief Calculate Drop Out Mask : x * 1.0/(1.0-rate)
+   * @param dropout drop out rate
+   * @retval Tensor& reference of drop out mask
+   */
+  TensorV2 dropout_mask(float dropout) const;
+
+  /**
+   * @brief Calculate Drop Out Mask : x * 1.0/(1.0-rate) inplace
+   * @param dropout drop out rate
+   */
+  void dropout_mask(float dropout);
+
+  /**
+   * @brief Calculate filter mask
+   * @param mask_len length of each mask along the last axis
+   * @param invert invert the mask
+   */
+  void filter_mask(const TensorV2 &mask_len, bool reverse = false);
+
+  /**
+   * @brief Calculate 2 Zone Out Mask
+   * @details Calculate zone out mask according to the bernoulli distribution.
+   * Zone out mask with rate @a zoneout for inplace and the other zone out mask
+   * with rate @a (1-zoneout).
+   * @param zoneout zone out rate
+   * @retval Tensor zone out mask for opposite tensor
+   */
+  TensorV2 zoneout_mask(float zoneout);
+
+  /**
+   * @brief Calculate 2 Zone Out Mask
+   * @details Calculate zone out mask according to the bernoulli distribution.
+   * Zone out mask with rate @a zoneout for inplace and the other zone out mask
+   * with rate @a (1-zoneout).
+   * @param opposite opposite zone out mask
+   * @param zoneout zone out rate
+   */
+  void zoneout_mask(TensorV2 &opposite, float zoneout);
+
   /**
    * @brief     Print element
    * @param[in] out out stream