Merge pull request #87 from lucasdelimanogueira/tmp

Tmp

README.md

@@ -29,7 +29,7 @@ x1 = norch.Tensor([[1, 2],
                   [3, 4]], requires_grad=True).to("cuda")
 
 x2 = norch.Tensor([[4, 3], 
-                  [2, 1]], requires_grad=True).to("cuda)
+                  [2, 1]], requires_grad=True).to("cuda")
 
 x3 = x1 @ x2
 result = x3.sum()

examples/train_multigpu.py

@@ -38,6 +38,7 @@ def main():
         ]
     )
 
+    print("Loading data")
     train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
     distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank)
     train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler)
@@ -58,6 +59,7 @@ def main():
             
             return out
 
+    print("Creating model")
     model = MyModel().to(device)
     model = DistributedDataParallel(model)
     criterion = nn.CrossEntropyLoss()
@@ -67,8 +69,15 @@ def main():
     print(f"Starting training on Rank {rank}/{world_size}\n\n")
 
     for epoch in range(epochs):    
+        
+        avg_loss = 0  
+        num_steps = 0  
+
         for idx, batch in enumerate(train_loader):
 
+            if idx % 100 == 0 and rank == 0:
+                print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}")
+
             inputs, target = batch
 
             inputs = inputs.to(device)
@@ -84,9 +93,13 @@ def main():
 
             optimizer.step()
 
+            avg_loss += loss[0]
+            num_steps += 1
+
+        avg_loss = avg_loss / num_steps
         if rank == 0:
-            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
-            loss_list.append(loss[0])
+            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')
+            loss_list.append(avg_loss)
 
 if __name__ == "__main__":
     main()

examples/train_nn_mnist.ipynb

@@ -139,7 +139,11 @@
     "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
     "loss_list = []\n",
     "\n",
-    "for epoch in range(epochs):    \n",
+    "for epoch in range(epochs):  \n",
+    "    \n",
+    "    avg_loss = 0  \n",
+    "    num_steps = 0\n",
+    "    \n",
     "    for idx, batch in enumerate(train_loader):\n",
     "\n",
     "        inputs, target = batch\n",
@@ -157,8 +161,12 @@
     "\n",
     "        optimizer.step()\n",
     "\n",
-    "    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')\n",
-    "    loss_list.append(loss[0])"
+    "        avg_loss += loss[0]\n",
+    "        num_steps += 1\n",
+    "\n",
+    "    avg_loss = avg_loss / num_steps\n",
+    "    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')\n",
+    "    loss_list.append(avg_loss)"
    ]
   },
   {

examples/train_singlegpu.py

@@ -24,6 +24,8 @@ def main():
         ]
     )
 
+
+    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)
 
@@ -43,14 +45,23 @@ def main():
             
             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 % 100 == 0:
+                print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}")
+
             inputs, target = batch
 
             inputs = inputs.to(device)
@@ -65,6 +76,13 @@ def main():
 
             optimizer.step()
         
+            avg_loss += loss[0]
+            num_steps += 1
+
+        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()

norch/csrc/cuda.cu

@@ -42,6 +42,10 @@ __host__ void cuda_to_cpu(Tensor* tensor) {
     strcpy(tensor->device, device_str); 
 }
 
+__host__ void free_cuda(float* data) {
+    cudaFree(data);
+}
+
 __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
     
     int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -136,7 +140,7 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
     __syncthreads();
 
     // Perform block-wise reduction
-    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) { // s>>=1 --> s = s/2
         if (tid < s) {
             partial_sum[tid] += partial_sum[tid + s];
         }

norch/csrc/cuda.h

@@ -4,6 +4,7 @@
 
     __host__ void cpu_to_cuda(Tensor* tensor, int device_id);
     __host__ void cuda_to_cpu(Tensor* tensor);
+    __host__ void free_cuda(float* data);
     
     __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);

norch/csrc/tensor.cpp

@@ -16,28 +16,33 @@ extern "C" {
             fprintf(stderr, "Memory allocation failed\n");
             exit(1);
         }
-        tensor->data = data;
-        tensor->shape = shape;
         tensor->ndim = ndim;
 
-        tensor->device = (char*)malloc(strlen(device) + 1);
-        if (device != NULL) {
-            strcpy(tensor->device, device);
-        } else {
-            fprintf(stderr, "Memory allocation failed\n");
-            exit(-1);
-        }
-
         tensor->size = 1;
         for (int i = 0; i < ndim; i++) {
             tensor->size *= shape[i];
         }
 
+        tensor->device = strdup(device);
         tensor->strides = (int*)malloc(ndim * sizeof(int));
-        if (tensor->strides == NULL) {
+        tensor->shape = (int*)malloc(ndim * sizeof(int));
+
+        if (tensor->device == NULL || tensor->strides == NULL || tensor->shape == NULL) {
             fprintf(stderr, "Memory allocation failed\n");
             exit(1);
         }
+
+        memcpy(tensor->shape, shape, tensor->ndim * sizeof(int));
+        strcpy(tensor->device, device);
+        
+        if (strcmp(tensor->device, "cpu") == 0) {
+            tensor->data = (float*)malloc(tensor->size * sizeof(float));
+            memcpy(tensor->data, data, tensor->size * sizeof(float));
+        } else {
+            //already allocated on gpu
+            tensor->data = data;
+        }
+            
         int stride = 1;
         for (int i = ndim - 1; i >= 0; i--) {
             tensor->strides[i] = stride;
@@ -47,6 +52,43 @@ extern "C" {
         return tensor;
     }
 
+    void delete_tensor(Tensor* tensor) {
+        if (tensor != NULL) {
+            
+            if (tensor->strides != NULL) {
+                free(tensor->strides);
+                tensor->strides = 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);
+        }
+    }
+
+
     float get_item(Tensor* tensor, int* indices) {
         int index = 0;
         for (int i = 0; i < tensor->ndim; i++) {

norch/csrc/tensor.h

@@ -5,8 +5,6 @@ typedef struct {
     float* data;
     int* strides;
     int* shape;
-    int* strides_cuda;
-    int* shape_cuda;
     int ndim;
     int size;
     char* device;
@@ -14,6 +12,7 @@ typedef struct {
 
 extern "C" {
     Tensor* create_tensor(float* data, int* shape, int ndim, char* device);
+    void delete_tensor(Tensor* tensor);
     float get_item(Tensor* tensor, int* indices);
     Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2);
     Tensor* sum_tensor(Tensor* tensor, int axis, bool keepdims);

norch/optim/optimizers/sgd.py

@@ -26,3 +26,4 @@ class SGD(Optimizer):
 
             parameter.detach()
             velocity.detach()
+            del parameter

norch/tensor.py

@@ -46,7 +46,6 @@ class Tensor:
             Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_char_p]
             Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor)
             
-
             self.tensor = Tensor._C.create_tensor(
                 self.data_ctype,
                 self.shape_ctype,
@@ -54,6 +53,10 @@ class Tensor:
                 self.device_ctype
             )
 
+            del self.data_ctype
+            del self.shape_ctype
+            del self.device_ctype
+        
         else:
             self.tensor = None,
             self.shape = None,
@@ -81,6 +84,12 @@ class Tensor:
         flat_data, shape = flatten_recursively(nested_list)
         return flat_data, shape
     
+    def __del__(self):
+        if self.tensor is not None:
+            Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
+            Tensor._C.delete_tensor.restype = None
+            Tensor._C.delete_tensor(self.tensor)
+    
     def __setattr__(self, name, value):
         if name == 'grad':
             for hook in self.hooks: