local machine version

source/sample/transformer/T2TTrainer.cpp

@@ -114,6 +114,7 @@ void T2TTrainer::Init(int argc, char ** argv)
     LoadParamBool(argc, argv, "epochcheckpoint", &useEpochCheckpoint, false);
     LoadParamInt(argc, argv, "updatestep", &updateStep, 1);
     LoadParamBool(argc, argv, "doubledend", &isDoubledEnd, false);
+    LoadParamBool(argc, argv, "smallbatch", &isSmallBatch, false);
 
     buf  = new int[bufSize];
     buf2 = new int[bufSize];
@@ -648,7 +649,8 @@ int T2TTrainer::LoadBatchLM(FILE * file,
         if(max < wn)
             max = wn;
 
-        if(sc >= sBatch && wc >= wBatch)
+        int tc = isSmallBatch ? max * sc : wc;
+        if(sc >= sBatch && tc >= wBatch)
             break;
     }
 
@@ -773,7 +775,8 @@ int T2TTrainer::LoadBatchMT(FILE * file,
         if(maxDec < wnDec)
             maxDec = wnDec;
 
-        if(sc >= sBatch * 2 && wcEnc >= wBatch)
+        int tc = isSmallBatch ? maxEnc * sc / 2 : wcEnc;
+        if(sc >= sBatch * 2 && tc >= wBatch)
             break;
     }
 

source/sample/transformer/T2TTrainer.h

@@ -133,6 +133,11 @@ public:
     
     /* indicates whether we double the </s> symble for the output of lms */
     bool isDoubledEnd;
+    
+    /* indicates whether we use batchsize = max * sc
+       rather rather than batchsize = word-number, where max is the maximum
+       length and sc is the sentence number */
+    bool isSmallBatch;
 
 public:
     /* constructor */