Merge pull request #69 from ivanallen/memleak

fix memleak
serizba · Nov 25, 2020 · 2d6f941 · 2d6f941
2 parents efabf78 + 5424d5f
commit 2d6f941
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 3 deletions.
diff --git a/examples/load_model/CMakeLists.txt b/examples/load_model/CMakeLists.txt
@@ -4,7 +4,9 @@ project(example)
 find_library(TENSORFLOW_LIB tensorflow HINT $ENV{HOME}/libtensorflow2/lib)
 
 set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-omit-frame-pointer -fsanitize=address")
+set(CMAKE_LINKER_FLAGS "${CMAKE_LINKER_FLAGS} -lasan")
 
 add_executable(example main.cpp)
 target_include_directories(example PRIVATE ../../include $ENV{HOME}/libtensorflow2/include)
-target_link_libraries (example "${TENSORFLOW_LIB}")
+target_link_libraries (example "${TENSORFLOW_LIB}")
diff --git a/examples/load_model/main.cpp b/examples/load_model/main.cpp
@@ -12,5 +12,10 @@ int main() {
 
     std::cout << output << std::endl;
 
+    auto values = output.get_data<float>();
+
+    for (auto v : values) {
+        std::cout << v << std::endl;
+    }
     return 0;
 }
diff --git a/include/cppflow/defer.h b/include/cppflow/defer.h
@@ -0,0 +1,25 @@
+#pragma once
+#include <functional>
+
+namespace cppflow {
+
+class defer {
+public:
+    typedef std::function<void ()> Func;
+
+    explicit defer(const Func& func) : _func(func) {}
+    ~defer() {
+        _func();
+    }
+
+    defer(const defer&) = delete;
+    defer(defer&&) = delete;
+    defer& operator=(const defer&) = delete;
+    void* operator new (size_t) = delete;
+    void operator delete (void*) = delete;
+
+private:
+    Func _func;
+};
+
+} // namespace cppflow
diff --git a/include/cppflow/model.h b/include/cppflow/model.h
@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "context.h"
+#include "defer.h"
 #include "tensor.h"
 
 namespace cppflow {
@@ -84,6 +85,12 @@ namespace cppflow {
 
         std::vector<TF_Output> inp_ops(inputs.size());
         std::vector<TF_Tensor*> inp_val(inputs.size());
+
+        defer d([&inp_val]{
+            for (auto* tf_tensor : inp_val) {
+                TF_DeleteTensor(tf_tensor);
+            }
+        });
         for (int i=0; i<inputs.size(); i++) {
 
             // Operations

diff --git a/include/cppflow/tensor.h b/include/cppflow/tensor.h
@@ -167,14 +167,16 @@ namespace cppflow {
 
         // EXECUTE
         int n = 1;
-        TFE_TensorHandle* res[1];
+        TFE_TensorHandle* res[1] = { nullptr };
         TFE_Execute(op, res, &n, context::get_status());
         status_check(context::get_status());
         TFE_DeleteOp(op);
 
         tensor r;
         r.tf_tensor = { TFE_TensorHandleResolve(res[0], context::get_status()), TF_DeleteTensor};
         status_check(context::get_status());
+        TFE_DeleteTensorHandle(res[0]);
+
         r.tfe_handle = {TFE_NewTensorHandle(r.tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};
         status_check(context::get_status());
 
@@ -205,7 +207,10 @@ namespace cppflow {
 
         // Convert to correct type
         const auto T_data = static_cast<T*>(raw_data);
-        return std::vector<T>(T_data, T_data + size);
+        std::vector<T> r(T_data, T_data + size);
+        TF_DeleteTensor(res_tensor);
+
+        return r;
     }
 
     datatype tensor::dtype() const {