Do not populate beam entries for blanks

amoliu · Jun 26, 2017 · 90a1463 · 90a1463
1 parent d6399bb
commit 90a1463
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/pytorch_ctc/src/ctc_beam_entry.h b/pytorch_ctc/src/ctc_beam_entry.h
@@ -55,23 +55,23 @@ struct BeamEntry {
   // create a vector of children.  The object pointed to by p
   // cannot be copied and should not be moved, otherwise parent will
   // become invalid.
-  BeamEntry(BeamEntry* p, int l, int L, int t) : parent(p), label(l) {
-    PopulateChildren(L);
+  BeamEntry(BeamEntry* p, int l, int L, int blank) : parent(p), label(l) {
+    PopulateChildren(L, blank);
   }
   inline bool Active() const { return newp.total != kLogZero; }
   inline bool HasChildren() const { return !children.empty(); }
-  void PopulateChildren(int L) {
+  void PopulateChildren(int L, int blank) {
     if (HasChildren()) {
       return;
     }
-    children = std::vector<BeamEntry>(L);
-    int ci = 0;
-    for (auto& c : children) {
+    children = std::vector<BeamEntry>(L-1);
+    for (int ci = 0; ci < L; ++ci) {
+      if (ci == blank) continue;
       // The current object cannot be copied, and should not be moved.
       // Otherwise the child's parent will become invalid.
+      auto& c = children[ci];
       c.parent = this;
       c.label = ci;
-      ++ci;
     }
   }
   inline std::vector<BeamEntry>* Children() {

diff --git a/pytorch_ctc/src/ctc_beam_search.h b/pytorch_ctc/src/ctc_beam_search.h
@@ -299,7 +299,7 @@ void CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::Step(
     }
 
     if (!b->HasChildren()) {
-      b->PopulateChildren(num_classes_);
+      b->PopulateChildren(num_classes_, blank_index_);
     }
 
     for (BeamEntry& c : *b->Children()) {
@@ -348,7 +348,7 @@ void CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::Reset() {
 
   // This beam root, and all of its children, will be in memory until
   // the next reset.
-  beam_root_.reset(new BeamEntry(nullptr, -1, num_classes_, -1));
+  beam_root_.reset(new BeamEntry(nullptr, -1, num_classes_, blank_index_));
   beam_root_->newp.total = 0.0;  // ln(1)
   beam_root_->newp.blank = 0.0;  // ln(1)