fix: reduce memory leak by 85%

csrc/layers/convolutional_layer.h

@@ -118,21 +118,21 @@ class Conv2D : public Layer {
       for (int oh = 0; oh < output_height; ++oh) {
         for (int ow = 0; ow < output_width; ++ow) {
           // Compute the dot product of the kernel and the input patch
-          double result = 0.0;
+          std::shared_ptr<Value> result = std::make_shared<Value>(0.0);
           for (int ic = 0; ic < in_channels; ++ic) {
             for (int kh = 0; kh < kernel_size; ++kh) {
               for (int kw = 0; kw < kernel_size; ++kw) {
                 int ih = oh * stride + kh - padding;
                 int iw = ow * stride + kw - padding;
                 if (ih >= 0 && ih < height && iw >= 0 && iw < width) {
-                  result += input->get({ic, ih, iw})->data *
-                      weights->get({oc, ic, kh, kw})->data;
+                  result = result->add(input->get({ic, ih, iw})->mul(
+                      weights->get({oc, ic, kh, kw})));
                 }
               }
             }
           }
-          result += bias->get(oc)->data; // Add bias
-          output->set({oc, oh, ow}, std::make_shared<Value>(result));
+          result = result->add(bias->get(oc)); // Add bias
+          output->set({oc, oh, ow}, result);
         }
       }
     }

csrc/main.cc

@@ -117,8 +117,8 @@ PYBIND11_MODULE(_core, m) {
       .def_readonly("minIdx", &Tensor::minIdx)
       .def_readonly("vals", &Tensor::v)
       .def("normalize_idx", &Tensor::normalize_idx)
-      .def("backward", &Tensor::backward)
       .def("zero_grad", &Tensor::zero_grad)
+      .def("reshape", &Tensor::reshape)
       .def("__add__", &Tensor::add)
       .def("__truediv__", &Tensor::div)
       .def("matmul", &Tensor::matmul)

csrc/tensor.h

@@ -23,6 +23,12 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     this->compute_stride();
   }
 
+  ~Tensor(){
+    this->strides.clear();
+    this->shape.clear();
+    this->v.clear();
+  }
+
   void compute_stride() {
     this->strides.clear();
     this->strides.resize(this->shape.size());
@@ -39,6 +45,24 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     this->maxIdx--; // 1 less
   }
 
+  void reshape(std::vector<int> new_shape){
+    int total_ele = 1;
+    for(auto &e:new_shape){
+      total_ele*=e;
+    }
+    if(total_ele!=this->maxIdx +1 ){
+      throw std::runtime_error("New shape must be able to contain ("
+        + std::to_string(this->maxIdx+1)
+        + "), but new shape can handle: "
+        + std::to_string(total_ele)
+        + " elements." 
+      );
+    }
+
+    this->shape = new_shape;
+    this->compute_stride();
+  }
+
   std::string tensor_shape_str() {
     std::string shape_str = "(";
     for (auto& e : this->shape) {
@@ -113,12 +137,6 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     }
   }
 
-  void backward() {
-    for (auto& e : this->v) {
-      e->backward();
-    }
-  }
-
   void remove_redundant_rows(std::shared_ptr<Tensor> t) {
     // remove redundant rows(1)
     std::vector<int> new_shape = {};

csrc/value.cc

@@ -46,6 +46,8 @@ void Value::backward() {
   for (int i = int(topo_list.size()) - 1; i >= 0; i--) {
     // std::cout << "i: " << i << "; node: " << topo_list[i]->printMe() << "\n";
     topo_list[i]->executeBackWardMethod();
+    topo_list[i]->clearBackwardMethod();
+    topo_list[i]->_prev.clear();
   }
 }
 

csrc/value.h

@@ -32,19 +32,28 @@ class Value : public std::enable_shared_from_this<Value> {
   Value(double data) : data(data) {}
   Value(double data, std::unordered_set<std::shared_ptr<Value>> _prev, char _op)
       : data(data), _prev(std::move(_prev)), _op(_op) {}
+
+  ~Value(){
+    this->_prev.clear();
+    this->clearBackwardMethod();
+  }
 
   // Setter to assign a new function
   void setBackWardMethod(std::function<void()> func) {
-    backward_ = func;
+    this->backward_ = func;
   }
 
-  // Method to execute the private method
+  // Method to execute the private method backward_
   void executeBackWardMethod() {
-    if (backward_) {
-      backward_();
+    if (this->backward_) {
+      this->backward_();
     }
   }
 
+  void clearBackwardMethod() {
+    this->backward_ = nullptr;
+  }
+
   void backward();
 
   std::string printMe() {