faster dropout implementation

source/sample/transformer/T2TEncoder.cpp

@@ -154,7 +154,7 @@ XTensor AttEncoder::Make(XTensor &input, XTensor &mask, bool skipInputRes, bool
         x = fnnLayerNorms[i].Make(res);
 
         if(isTraining && dropoutP > 0)
-            x = Dropout(x);
+            x = Dropout(x, dropoutP);
     }
 
     return x;

source/sample/transformer/T2TTrainer.cpp

@@ -166,16 +166,6 @@ void T2TTrainer::Train(const char * fn, T2TModel * model)
             
             /* get probabilities */
             float prob = GetProb(&output, &gold, NULL);
-
-            MTYPE totalUsed = 0;
-            MTYPE totalSize = 0;
-            
-            for (int i = 0; i <= mem->curBlockID; i++) {
-                totalSize += mem->blocks[i].size;
-                totalUsed += mem->blocks[i].used;
-            }
-
-            //fprintf(stderr, "%d(%ld,%ld,%f)\n", mem->curBlockID, totalUsed, totalSize, (float)totalUsed/totalSize);
             
             loss += -prob;
             wordCount += wc;
@@ -209,6 +199,8 @@ void T2TTrainer::Train(const char * fn, T2TModel * model)
 
     fclose(tf);
     
+    epoch = MIN(epoch, nepoch);
+    
     XPRINT6(0, stderr, "[INFO] lr=%.2e, elapsed=%.1fs, step=%d, epoch=%d, word=%d, ppl=%.3f\n",
             lr, elapsed, step, epoch, wordCountTotal, exp(loss / wordCount));
     XPRINT3(0, stderr, "[INFO] training finished (took %.1fs, step=%d and epoch=%d)\n",

source/tensor/function/Dropout.cpp

@@ -25,6 +25,7 @@
 #include "Dropout.h"
 #include "Dropout.cuh"
 #include "../core/arithmetic/Multiply.h"
+#include "../core/arithmetic/SumDim.h"
 #include "../core/math/ScaleAndShift.h"
 
 namespace nts{ // namespace nts(NiuTrans.Tensor
@@ -148,13 +149,15 @@ the same inference procedure as that with no use of dropout on the test data.
 >> x - input tensor
 >> y - output tensor
 >> prob - probability to set an element to zero
+>> leadDim - the dimension along which we generate the random numbers
 */
-XTensor Dropout(const XTensor &x, DTYPE prob)
+XTensor Dropout(const XTensor &x, DTYPE prob, int leadDim)
 {
+    int n = leadDim < 0 ? x.order - 1 : leadDim;
     DTYPE scaleFactor = (DTYPE)1.0 / ((DTYPE)1.0 - prob);
     
     /* generate a mask tensor with probability p */
-    int unitNum = x.unitNum;
+    int unitNum = x.dimSize[n];
     DTYPE * maskArray = new DTYPE[unitNum];
 
     srand((unsigned int)time(NULL));
@@ -162,9 +165,15 @@ XTensor Dropout(const XTensor &x, DTYPE prob)
         maskArray[i] = RandomBernoulli(prob, scaleFactor);
     
     XTensor mask(&x);
-    mask.SetData(maskArray, unitNum);
+    mask.SetZeroAll();
+    
+    XTensor * maskVector = NewTensorBuf(1, &unitNum, X_FLOAT, 1.0F, x.devID, x.mem);
+    maskVector->SetData(maskArray, unitNum);
+    
+    _SumDim(&mask, maskVector, &mask, n);
 
     delete[] maskArray;
+    DelTensorBuf(maskVector);
     
     return Multiply(x, mask);
 }

source/tensor/function/Dropout.h

@@ -30,19 +30,19 @@ namespace nts{ // namespace nts(NiuTrans.Tensor)
 /* generate a random bernoulli number */
 inline DTYPE RandomBernoulli(DTYPE prob, DTYPE value)
 {
-    return (DTYPE)rand()/(DTYPE)RAND_MAX > prob ? (DTYPE)value : (DTYPE)0.0;
+    return (DTYPE)rand()/(DTYPE)RAND_MAX >= prob ? (DTYPE)value : 0;
 }
 
 /* dropout function */
-void _Dropout(const XTensor * x, XTensor * y, unsigned int seed, DTYPE prob = 0.5);
+void _Dropout(const XTensor * x, XTensor * y, unsigned int seed, DTYPE prob);
 
 /* de/dx */
 void _DropoutBackward(const XTensor * y, const XTensor * x, 
                       const XTensor * dedy, XTensor * dedx, 
-                      unsigned int seed, DTYPE prob = 0.5);
+                      unsigned int seed, DTYPE prob);
     
 /* dropout function */
-XTensor Dropout(const XTensor &x, DTYPE prob = 0.5);
+XTensor Dropout(const XTensor &x, DTYPE prob, int leadDim = -1);
 
 } // namespace nts(NiuTrans.Tensor)
 

source/tensor/test/TDropout.cpp

@@ -56,7 +56,7 @@ bool TestDropout1()
     float prob = 0.2F;
     int seed = 20;
     _Dropout(x, y, seed, prob);
-    yUser = Dropout(*x);
+    yUser = Dropout(*x, 0.5F);
 
     /* check result */
     int zeroNum1 = 0;
@@ -92,7 +92,7 @@ bool TestDropout1()
 
     /* call Dropout function */
     _Dropout(xGPU, yGPU, seed, prob);
-    yUserGPU = Dropout(*xGPU);
+    yUserGPU = Dropout(*xGPU, 0.5F);
 
     /* check result */
     zeroNum1 = 0;