Merge pull request parlance#30 from joshemorris/mem-leak

Fix parlance#14 Memory Leaks
amoliu · Oct 28, 2017 · d2e728c · d2e728c
2 parents 83a001d + f7ad6d0
commit d2e728c
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 1 deletion.
diff --git a/pytorch_ctc/__init__.py b/pytorch_ctc/__init__.py
@@ -81,6 +81,11 @@ def __init__(self, labels, lm_path, trie_path, blank_index=0, space_index=28):
         self._scorer = ctc._get_kenlm_scorer(labels, len(labels), space_index, blank_index, lm_path.encode(),
                                              trie_path.encode())
 
+    # This is a way to make sure the destructor is called for the C++ object
+    # Frees all the member data items that have allocated memory
+    def __del__(self):
+        ctc._free_kenlm_scorer(self._scorer)
+
     def set_lm_weight(self, weight):
         if weight is not None:
             ctc._set_kenlm_scorer_lm_weight(self._scorer, weight)

diff --git a/pytorch_ctc/src/cpu_binding.cpp b/pytorch_ctc/src/cpu_binding.cpp
@@ -71,6 +71,7 @@ namespace pytorch {
     root.WriteToStream(ofs);
     ifs.close();
     ofs.close();
+    delete model;
     return 0;
   }
   #endif
@@ -88,6 +89,11 @@ namespace pytorch {
       #endif
     }
 
+    void free_kenlm_scorer(void* kenlm_scorer) {
+      ctc::KenLMBeamScorer* beam_scorer = static_cast<ctc::KenLMBeamScorer*>(kenlm_scorer);
+      delete beam_scorer;
+    }
+
     void set_kenlm_scorer_lm_weight(void *scorer, float weight) {
       #ifdef INCLUDE_KENLM
       ctc::KenLMBeamScorer *beam_scorer = static_cast<ctc::KenLMBeamScorer *>(scorer);

diff --git a/pytorch_ctc/src/cpu_binding.h b/pytorch_ctc/src/cpu_binding.h
@@ -8,6 +8,8 @@ typedef enum {
 int kenlm_enabled();
 void* get_kenlm_scorer(const wchar_t* label_str, int labels_size, int space_index, int blank_index,
                        const char* lm_path, const char* trie_path);
+void free_kenlm_scorer(void* kenlm_scorer);
+
 void set_kenlm_scorer_lm_weight(void *scorer, float weight);
 void set_kenlm_scorer_wc_weight(void *scorer, float weight);
 void set_kenlm_scorer_vwc_weight(void *scorer, float weight);

diff --git a/pytorch_ctc/src/ctc_beam_scorer_klm.h b/pytorch_ctc/src/ctc_beam_scorer_klm.h
@@ -48,9 +48,10 @@ namespace ctc_beam_search {
   class KenLMBeamScorer : public BaseBeamScorer<KenLMBeamState> {
    public:
 
-    virtual ~KenLMBeamScorer() {
+    ~KenLMBeamScorer() {
       delete model;
       delete trieRoot;
+      delete labels;
     }
     KenLMBeamScorer(Labels *labels, const char *kenlm_path, const char *trie_path)
                       : lm_weight(1.0f),