Fix recursion backward increasing a lot

norch/nn/module.py

@@ -51,9 +51,11 @@ class Module(ABC):
             parameter.zero_grad()
 
     def to(self, device):
-        for parameter in self.parameters():
+        for _, _, parameter in self.parameters():
             parameter.to(device)
 
+        return self
+    
     def state_dict(self):
         state = OrderedDict()
         for i, param in enumerate(self.parameters()):

norch/tensor.py

@@ -137,27 +137,36 @@ class Tensor:
     def backward(self, gradient=None):
         if not self.requires_grad:
             return
-        
+                
         if gradient is None:
             if self.shape == [1]:
                 gradient = Tensor([1])
             else:
                 raise RuntimeError("Gradient argument must be specified for non-scalar tensors.")
 
-        if self.grad is None:
-            self.grad = gradient
+        stack = [(self, gradient)]
+        visited = set()
+    
+        while stack:
+            tensor, grad = stack.pop()
+            
+            if tensor.grad is None:
+                tensor.grad = grad
+            else:
+                tensor.grad += grad
 
-        else:
-            self.grad += gradient
-
-        if self.grad_fn is not None: # not a leaf
-            grads = self.grad_fn.backward(gradient)
-            for tensor, grad in zip(self.grad_fn.input, grads):
-                if isinstance(tensor, Tensor):
-                    tensor.backward(grad)
+            # Propagate gradients to inputs if not a leaf tensor
+            if tensor.grad_fn is not None:
+                grads = tensor.grad_fn.backward(grad)
+                for tensor, grad in zip(tensor.grad_fn.input, grads):
+                    if isinstance(tensor, Tensor) and tensor not in visited:
+                        stack.append((tensor, grad))
+                        visited.add(tensor)
 
     def zero_grad(self):
-        self.grad = None
+        tmp = self.zeros_like()
+        self.detach()
+        self.grad = tmp
 
     def __getitem__(self, indices):
         if len(indices) != self.ndim:
@@ -570,5 +579,13 @@ class Tensor:
         result_data.shape = self.shape.copy()[::-1]
         result_data.ndim = self.ndim
         result_data.device = self.device
+        
+        result_data.requires_grad = self.requires_grad
 
-        return result_data
+        return result_data
+    
+    def detach(self):
+        self.grad = None
+        self.grad_fn = None
+
+        return self

test.py

@@ -19,6 +19,7 @@ if __name__ == "__main__":
     import time
     import random
     import numpy as np
+    import psutil
 
     """a = norch.Tensor([
         [[1.234, 2.123], [3.635, 4.456], [5.678, 6.789]],
@@ -69,11 +70,17 @@ if __name__ == "__main__":
 
     import norch.nn as nn
 
+    cpu_percent = psutil.cpu_percent(interval=1)
+    print(f"CPU Usage: {cpu_percent}%")
+    memory_usage = psutil.virtual_memory()
+    print(f"Memory Usage: {memory_usage.percent}%")  
+    
+
     class MeuModulo(nn.Module):
         def __init__(self):
             super(MeuModulo, self).__init__()
 
-            self.layer1 = nn.Linear(10, 2)
+            self.layer1 = nn.Linear(5, 2)
             #self.layer2 = nn.Linear(5, 2)
             self.sigmoid = nn.Sigmoid()
 
@@ -86,33 +93,26 @@ if __name__ == "__main__":
             return out
         
     modelo = MeuModulo()
-    input_list = [[0.05 for _ in range(10)]]
+    input_list = [[0.05 for _ in range(5)]]
     input = norch.Tensor(input_list).T
     criterion = nn.MSELoss()
-    optimizer = norch.optim.SGD(modelo.parameters(), lr=0.5)
+    optimizer = norch.optim.SGD(modelo.parameters(), lr=0.1)
 
     target_list = [[0.1 for _ in range(2)]]
     target = norch.Tensor(target_list).T
     
-    for epoch in range(5):
+    for epoch in range(50):
+        cpu_percent = psutil.cpu_percent(interval=1)
+        memory_usage = psutil.virtual_memory()
         output = modelo(input)
         loss = criterion(output, target)
         optimizer.zero_grad()
-        #print('FORA ANTES: ', modelo.layer1.bias, '\n')
 
-        #print(modelo.layer1.weight.grad)        
         loss.backward()
-        #print(modelo.layer1.weight.grad)
         optimizer.step()
-        #print('FORA DEPOIS: ', modelo.layer1.bias, '\n')
 
-        #print("\n\n")
-        #print(modelo.layer1.weight.grad)
-
-        #print(modelo.layer1.weight[0,0])
         print(loss)
 
-    exit()
 
     #### testar transpose axes!!!! make it contiguous
 
@@ -122,15 +122,20 @@ if __name__ == "__main__":
                         [[19, 20], [21, 22], [23, 24]],
                         [[25, 26], [27, 28], [29, 0.030]]], requires_grad=True)
     
-    op = nn.Sigmoid()
-    result = op(tensor1)
-    result = result.sum()
+    #op = nn.Sigmoid()
+    tensor2 = 
+    result = tensor1.sum()
+    
     result.backward()
     print(tensor1.grad)
     exit()"""
 
     # Reshape tensor1 to 2x3x5
-    reshaped_tensor = tensor1.transpose(1, 0)
+    """tensor1 = norch.Tensor([[[1, 2], [3, 4], [5, 6]],
+                        [[7, 8], [9, 10], [11, 12]],
+                        [[13, 14], [15, 16], [17, 18]],
+                        [[19, 20], [21, 22], [23, 24]],
+                        [[25, 26], [27, 28], [29, 0.030]]], requires_grad=True)
 
     # Create a 5x4 tensor
     tensor2 = norch.Tensor([[1, 2, 3],
@@ -138,16 +143,14 @@ if __name__ == "__main__":
                             [9, 10, 11],
                             [13, 14, 15],
                             [17, 18, 19]])
-
-    tensor2 = tensor2.transpose(1,0)
     
 
     # Multiply reshaped_tensor by tensor2
-    result = tensor2 @ reshaped_tensor
+    result = tensor2 @ tensor1
 
     result = result.sum()
     result.backward()
-    print(tensor1.grad)
+    print(tensor1.grad)"""
 
     #print(a.shape, b.shape, result.shape)
     #c = result.sum()