fix the SubNestedSequenceLayer implementations.

paddle/gserver/layers/SubNestedSequenceLayer.cpp

@@ -31,16 +31,22 @@ public:
   void backward(const UpdateCallback& callback = nullptr) override;
 
 private:
-  void calSelectedCols(const MatrixPtr scores,
-                       const int* seqStartPos,
-                       const int* subSeqStartPos);
+  void reorganizeSeqInfo(const ICpuGpuVectorPtr seqStartPos,
+                         const ICpuGpuVectorPtr subSeqStartPos);
+  void calSelectedCols(const MatrixPtr selectedIndices,
+                       const std::vector<std::vector<int>> inputSeqInfo);
   void buildOutputSeqInfo();
 
   std::vector<int> outSeqStartInfo_;
   std::vector<int> outSubSeqStartInfo_;
 
-  MatrixPtr scoreOverInputSeq_;
+  // if the second input of this layer is on GPU memory, copy it to CPU memory.
+  MatrixPtr selIdsCpu_;
+  // reorganize sequenceStartPositions and subSequenceStartPositions altogether
+  // into a 2d vector to facilitate the sequence selection process.
+  std::vector<std::vector<int>> inputSeqInfo_;
 
+  // the final seleted row indices in a batch,
   // rowIdx_ and selectedRows_ actually share a same memory.
   IVectorPtr rowIndice_;
   std::vector<int> selectedRows_;
@@ -57,12 +63,47 @@ bool SubNestedSequenceLayer::init(const LayerMap& layerMap,
   return true;
 }
 
-void SubNestedSequenceLayer::calSelectedCols(const MatrixPtr selected_indices,
-                                             const int* seqStartPos,
-                                             const int* subSeqStartPos) {
+void SubNestedSequenceLayer::reorganizeSeqInfo(
+    const ICpuGpuVectorPtr seqStartPos, const ICpuGpuVectorPtr subSeqStartPos) {
+  int* seqStarts = seqStartPos->getMutableData(false);
+  int* subSeqStarts = subSeqStartPos->getMutableData(false);
+
+  int seqNum = seqStartPos->getSize() - 1;
+  inputSeqInfo_.resize(seqNum, std::vector<int>());
+  int seqIdx = 0;
+  for (size_t i = 0; i < subSeqStartPos->getSize(); ++i) {
+    inputSeqInfo_[seqIdx].push_back(subSeqStarts[i]);
+    if (subSeqStarts[i] == seqStarts[seqIdx + 1]) {
+      seqIdx++;
+      if (seqIdx == seqNum) return;
+      inputSeqInfo_[seqIdx].push_back(subSeqStarts[i]);
+    }
+  }
+}
+
+void SubNestedSequenceLayer::calSelectedCols(
+    const MatrixPtr selectedIndices,
+    const std::vector<std::vector<int>> inputSeqInfo) {
   selectedRows_.clear();
   outSubSeqStartInfo_.resize(1, 0);
   outSeqStartInfo_.resize(1, 0);
+
+  size_t seqNum = selectedIndices->getHeight();
+  size_t beamSize = selectedIndices->getWidth();
+  for (size_t i = 0; i < seqNum; ++i) {
+    for (size_t j = 0; j < beamSize; ++j) {
+      if (selectedIndices->getElement(i, j) == -1.) break;
+      int selSubSeqIdx = selectedIndices->getElement(i, j);
+      CHECK_GT(inputSeqInfo_[i].size() - 1, selSubSeqIdx);
+
+      size_t subSeqLen =
+          inputSeqInfo_[i][selSubSeqIdx + 1] - inputSeqInfo_[i][selSubSeqIdx];
+      for (size_t k = 0; k < subSeqLen; ++k)
+        selectedRows_.push_back(inputSeqInfo_[i][selSubSeqIdx] + k);
+      outSubSeqStartInfo_.push_back(outSubSeqStartInfo_.back() + subSeqLen);
+    }
+    outSeqStartInfo_.push_back(outSubSeqStartInfo_.back());
+  }
 }
 
 void SubNestedSequenceLayer::buildOutputSeqInfo() {
@@ -83,17 +124,35 @@ void SubNestedSequenceLayer::forward(PassType passType) {
   Layer::forward(passType);
 
   const Argument& inputSeq = getInput(0);
-  const MatrixPtr selected_indices = getInputValue(1);
   CHECK(inputSeq.hasSubseq()) << "The first input of SubNestSequence layer "
                               << "must be a nested sequence.";
-  CHECK_EQ(inputSeq.getNumSequences(), selected_indices->getHeight());
-
-  calSelectedCols(selected_indices,
-                  inputSeq.sequenceStartPositions->getMutableData(false),
-                  inputSeq.subSequenceStartPositions->getMutableData(false));
+  const MatrixPtr selectedIndices = getInputValue(1);
+  CHECK_EQ(inputSeq.getNumSequences(), selectedIndices->getHeight());
+
+  if (dynamic_cast<GpuMatrix*>(selectedIndices.get())) {
+    /*
+     * Currently, the second input for this layer generated by
+     * kmax_sequence_score_layer whose output is always stored on CPU,
+     * or a data_layer which canbe on GPU.
+     *
+     * If the second input is on GPU, copy it to CPU memory, because this
+     * input always uses very few memory, and operations related to it are
+     * all logic control, not computations.
+     */
+    Matrix::resizeOrCreate(selIdsCpu_,
+                           selectedIndices->getHeight(),
+                           selectedIndices->getWidth(),
+                           false /* trans */,
+                           false /* useGpu */);
+    selIdsCpu_->copyFrom(*selectedIndices);
+  } else {
+    selIdsCpu_ = selectedIndices;
+  }
 
+  reorganizeSeqInfo(inputSeq.sequenceStartPositions,
+                    inputSeq.subSequenceStartPositions);
+  calSelectedCols(selIdsCpu_, inputSeqInfo_);
   resetOutput(selectedRows_.size(), getSize());
-  buildOutputSeqInfo();
 
   if (useGpu_) {
     rowIndice_ = IVector::create(selectedRows_.size(), useGpu_);
@@ -103,6 +162,7 @@ void SubNestedSequenceLayer::forward(PassType passType) {
         IVector::create(selectedRows_.data(), selectedRows_.size(), useGpu_);
   }
 
+  buildOutputSeqInfo();
   getOutputValue()->selectRows(*getInputValue(0), *rowIndice_);
 }
 

python/paddle/trainer_config_helpers/layers.py

@@ -6097,9 +6097,11 @@ def sub_nested_seq_layer(input, selected_indices, name=None):
     The sub_nested_seq_layer accepts two inputs: the first one is a nested
     sequence; the second one is a set of selceted indices in the nested sequence.
 
+    Then sub_nest_seq_layer trims the first nested sequence input according to
+    the selected indices to form a new output.
+
+    This layer is useful in beam training.
 
-    Then sub_nest_seq_layer selects trims the first input according to the
-    selected indices to give a new output. This layer is used in beam training.
 
     The example usage is: