fix memory usage

norch/csrc/tensor.cpp

@@ -49,37 +49,30 @@ extern "C" {
 
     void delete_tensor(Tensor* tensor) {
         if (tensor != NULL) {
-            
-            if (tensor->shape != NULL) {
-                free(tensor->shape);
-                tensor->shape = NULL;
-            }
-            
-            if (tensor->data != NULL) {
-                if (strcmp(tensor->device, "cpu") == 0) {
-                    free(tensor->data);
-                } else {
-                    free_cuda(tensor->data);
-                }
-                tensor->data = NULL;
-            }
-
-            if (tensor->device != NULL) {
-                free(tensor->device);
-                tensor->device = NULL;
-            }
-
-            if (tensor->strides != NULL) {
-                free(tensor->strides);
-                tensor->strides = NULL;
-            }
-            
             free(tensor);
+            tensor = NULL;
+        }
+    }
+
+    void delete_shape(Tensor* tensor) {
+        if (tensor->shape != NULL) {
+            free(tensor->shape);
+            tensor->shape = NULL;
+        }
+    }
+
+    void delete_data(Tensor* tensor) {
+        if (tensor->data != NULL) {
+            if (strcmp(tensor->device, "cpu") == 0) {
+                free(tensor->data);
+            } else {
+                free_cuda(tensor->data);
+            }
+            tensor->data = NULL;
         }
     }
 
     void delete_strides(Tensor* tensor) {
-        // as strides always are allocated within C code, it must be handled separatedly
         if (tensor->strides != NULL) {
             free(tensor->strides);
             tensor->strides = NULL;
@@ -87,7 +80,6 @@ extern "C" {
     }
 
     void delete_device(Tensor* tensor) {
-        // as device always are allocated within C code, it must be handled separatedly
         if (tensor->device != NULL) {
             free(tensor->device);
             tensor->device = NULL;
@@ -114,20 +106,28 @@ extern "C" {
         int device_id = 0;
         char* endptr;
 
+        char* target_device_type;
+
         long num = strtol(target_device, &endptr, 10);
         if (*endptr == '\0') {
             device_id = (int)num;
-            target_device = new char[strlen("cuda") + 1];
+            target_device_type = new char[strlen("cuda") + 1];
-            strcpy(target_device, "cuda");
+            strcpy(target_device_type, "cuda");
+        }
+        else {
+            target_device_type = new char[strlen("cuda") + 1];
+            strcpy(target_device_type, "cpu");
         }
 
-        if ((strcmp(target_device, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
+        if ((strcmp(target_device_type, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
             cpu_to_cuda(tensor, device_id);
         }
 
-        else if ((strcmp(target_device, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {
+        else if ((strcmp(target_device_type, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {
             cuda_to_cpu(tensor);
         }
+
+        free(target_device_type);
     }
 
     Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2) {

norch/csrc/tensor.h

@@ -14,6 +14,8 @@ extern "C" {
     Tensor* create_tensor(float* data, int* shape, int ndim, char* device);
     void delete_tensor(Tensor* tensor);
     void delete_strides(Tensor* tensor);
+    void delete_shape(Tensor* tensor);
+    void delete_strides(Tensor* tensor);
     void delete_device(Tensor* tensor);
     float get_item(Tensor* tensor, int* indices);
     Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2);

norch/tensor.py

@@ -83,8 +83,7 @@ class Tensor:
     def __del__(self):
             
         if hasattr(self, '_data_ctype') and self._data_ctype is not None:
-            # tensor created by user (ctypes) will be deleted by python garbage collector
+        
-            # only strides need to be deallocated manually because it is created inside C code
             Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
             Tensor._C.delete_strides.restype = None
             Tensor._C.delete_strides(self.tensor)
@@ -93,8 +92,26 @@ class Tensor:
             Tensor._C.delete_device.restype = None
             Tensor._C.delete_device(self.tensor)
 
+            Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_tensor.restype = None
+            Tensor._C.delete_tensor(self.tensor)
+
         elif self.tensor is not None:
-            # tensor created during operations must be deallocated
+            Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_strides.restype = None
+            Tensor._C.delete_strides(self.tensor)
+
+            Tensor._C.delete_data.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_data.restype = None
+            Tensor._C.delete_data(self.tensor)
+
+            Tensor._C.delete_shape.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_shape.restype = None
+            Tensor._C.delete_shape(self.tensor)
+
+            Tensor._C.delete_device.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_device.restype = None
+            Tensor._C.delete_device(self.tensor)
 
             Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
             Tensor._C.delete_tensor.restype = None

train_singlegpu.py

@@ -0,0 +1,95 @@
+import norch
+import norch.nn as nn
+import norch.optim as optim
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+from memory_profiler import profile
+
+@profile
+def main():
+
+    BATCH_SIZE = 32
+    device = "cpu"
+    epochs = 10
+
+    transform = T.Compose(
+        [
+            T.ToTensor(),
+            T.Reshape([-1, 784, 1])
+        ]
+    )
+
+    target_transform = T.Compose(
+        [
+            T.ToTensor()
+        ]
+    )
+
+
+    print("Loading data")
+    train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+    train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE)
+
+    class MyModel(nn.Module):
+        def __init__(self):
+            super(MyModel, self).__init__()
+            self.fc1 = nn.Linear(784, 30)
+            self.sigmoid1 = nn.Sigmoid()
+            self.fc2 = nn.Linear(30, 10)
+            self.sigmoid2 = nn.Sigmoid()
+
+        def forward(self, x):
+            out = self.fc1(x)
+            out = self.sigmoid1(out)
+            out = self.fc2(out)
+            out = self.sigmoid2(out)
+            
+            return out
+
+    print("Creating model")
+    model = MyModel().to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+    loss_list = []
+
+    print("Starting training")
+    for epoch in range(epochs): 
+
+        avg_loss = 0  
+        num_steps = 0  
+        
+        for idx, batch in enumerate(train_loader):
+
+            if idx % 300 == 0 and idx > 0:
+                print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}")
+                break
+
+            inputs, target = batch
+
+            inputs = inputs.to(device)
+            target = target.to(device)
+
+            outputs = model(inputs)
+            loss = criterion(outputs, target)
+            
+            optimizer.zero_grad()
+            
+            loss.backward()
+
+            optimizer.step()
+        
+            avg_loss += loss[0]
+            num_steps += 1
+
+        break
+        avg_loss = avg_loss / num_steps
+        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')
+        loss_list.append(avg_loss)
+        
+
+
+if __name__ == "__main__":
+    main()
+