Merge pull request #1 from AshtonO/fix_predict_memory_leak

Fix a memory leak in session.run

ext/sciruby/tensorflow_c/files/tf_tensor_helper.cc

@@ -16,14 +16,18 @@
 
 namespace tensorflow {
 
+// Callback function for deallocating the data buffer of a tensor when
+// TF_DeleteTensor is called.
+void tensor_deallocator(void* data, size_t len, void* arg){
+    free(data);
+}
+
 TF_Tensor* TF_NewTensor_wrapper(TF_DataType dtype, long long* dims, int num_dims,
                                 void* data, size_t len) {
-
         void* cData = malloc(len);
         memcpy(cData, data, len);
 
-        return TF_NewTensor(dtype, dims, num_dims, cData, len, [] (void *cData, size_t len, void* arg){
-                            }, nullptr);
+        return TF_NewTensor(dtype, dims, num_dims, cData, len, &tensor_deallocator, nullptr);
 };
 
 long long tensor_size(TF_Tensor* tensor)
@@ -201,11 +205,19 @@ std::vector<TF_Tensor *> Session_run(TF_Session* graph_session, std::vector<TF_O
 
         TF_SessionRun(graph_session, NULL, &inputPorts_array[0], &inputValues_array[0], inputPorts.size(), &outputPorts_array[0], &outputValues_array[0], outputPorts.size(), &cTargets_array[0], cTargets.size(), NULL, status);
 
+
         array_length = outputPorts.size();
         std::vector<TF_Tensor *> TF_Tensor_vector;
         for (auto i = 0; i < array_length; ++i)
                 TF_Tensor_vector.push_back(outputValues_array[i]);
 
+        delete[] inputPorts_array;
+        delete[] outputPorts_array;
+        delete[] inputValues_array;
+        delete[] outputValues_array;
+        delete[] cTargets_array;
+        TF_DeleteStatus(status);
+
         return TF_Tensor_vector;
 };
 

lib/tensorflow/session.rb

@@ -33,6 +33,9 @@ def initialize(graph, c_options = nil)
     # On success, returns the fetched Tensors in the same order as supplied in
     # the fetches argument. If fetches is set to nil, the returned Tensor fetches
     # is empty.
+    #
+    # Note that the caller maintains responsibility for the input tensors, so 
+    # the caller should still call tensor.delete() on each input tensor
     def run(inputs, outputs, targets)
         inputPorts = Tensorflow::TF_Output_vector.new
         inputValues = Tensorflow::Tensor_Vector.new
@@ -57,6 +60,9 @@ def run(inputs, outputs, targets)
         outputValues.each do |value|
             converted_value = convert_value_for_output_array(value)
             output_array.push(converted_value)
+            # Since we're returning the results as an array, there's no point keeping
+            # the output tensor, so we delete it. 
+            Tensorflow::TF_DeleteTensor(value)
         end
         output_array
     end

lib/tensorflow/tensor.rb

@@ -60,6 +60,13 @@ def initialize(value, type = nil)
                                                        dimension_data, rank, tensor_data, data_size * flatten.length)
     end
 
+    # Deletes the tensor which will also ensure that the tensor_deallocator function is called
+    # and the C++ memory is freed properly. A memory leak will occur anytime the ruby garbage collector
+    # deletes a tensor before this function is called.
+    def delete()
+        Tensorflow::TF_DeleteTensor(@tensor)
+    end
+
     #
     # Helper function to automatically set the data type of tensor.
     #