sigmoid kernel avoid overflow

norch/autograd/functions.py

@@ -28,7 +28,7 @@ class AddBroadcastedBackward:
         for i in range(len(shape)):
             if shape[i] == 1:
                 gradient = gradient.sum(axis=i, keepdim=True)
-
+        
         return gradient
     
 class SubBackward:
@@ -183,6 +183,7 @@ class DivisionBackward:
         x, y = self.input
         grad_x = gradient / y
         grad_y = -1 * gradient * (x / (y * y))
+
         return [grad_x, grad_y]
     
 class SinBackward:
@@ -291,4 +292,13 @@ class CrossEntropyLossBackward:
 
         return [grad_logits, None]  # targets do not have a gradient
 
+class SigmoidBackward:
+    def __init__(self, input):
+        self.input = [input]
+
+    def backward(self, gradient):
+        sigmoid_x = self.input[0].sigmoid()
+        grad_input = gradient * sigmoid_x * (1 - sigmoid_x)
+
+        return [grad_input]
 

norch/csrc/cpu.cpp

@@ -405,6 +405,22 @@ void sin_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
+void sigmoid_tensor_cpu(Tensor* tensor, float* result_data) {
+    for (int i = 0; i < tensor->size; i++) {
+        // avoid overflow
+        if (tensor->data[i] >= 0) {
+
+            float z = expf(-tensor->data[i]);
+            result_data[i] = 1 / (1 + z);
+
+        } else {
+
+            float z = expf(tensor->data[i]);
+            result_data[i] = z / (1 + z);
+        }
+    }
+}
+
 void cos_tensor_cpu(Tensor* tensor, float* result_data) {
     for (int i = 0; i < tensor->size; i++) {
         result_data[i] = cosf(tensor->data[i]);

norch/csrc/cpu.h

@@ -32,5 +32,6 @@ void transpose_axes_cpu(Tensor* tensor, float* result_data, int axis1, int axis2
 void assign_tensor_cpu(Tensor* tensor, float* result_data);
 void sin_tensor_cpu(Tensor* tensor, float* result_data);
 void cos_tensor_cpu(Tensor* tensor, float* result_data);
+void sigmoid_tensor_cpu(Tensor* tensor, float* result_data);
 
 #endif /* CPU_H */

norch/csrc/cuda.cu

@@ -764,6 +764,39 @@ __host__ void cos_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void sigmoid_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (i < size) {
+        // avoid overflow
+        if (data[i] >= 0) {
+
+            float z = expf(-data[i]);
+            result_data[i] = 1 / (1 + z);
+
+        } else {
+
+            float z = expf(data[i]);
+            result_data[i] = z / (1 + z);
+        }
+    }
+}
+
+__host__ void sigmoid_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    sigmoid_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
 
 
 

norch/csrc/cuda.h

@@ -73,6 +73,9 @@
      __global__ void cos_tensor_cuda_kernel(float* data, float* result_data, int size);
     __host__ void cos_tensor_cuda(Tensor* tensor, float* result_data);
 
+    __global__ void sigmoid_tensor_cuda_kernel(float* data, float* result_data, int size);
+    __host__ void sigmoid_tensor_cuda(Tensor* tensor, float* result_data);
+
 
 
 

norch/csrc/tensor.cpp

@@ -1290,6 +1290,43 @@ extern "C" {
         }
     }
 
+    Tensor* sigmoid_tensor(Tensor* tensor) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            sigmoid_tensor_cuda(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            sigmoid_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
     Tensor* transpose_tensor(Tensor* tensor) {
         char* device = (char*)malloc(strlen(tensor->device) + 1);
         if (device != NULL) {

norch/nn/functional.py

@@ -1,9 +1,11 @@
 import math
 import norch
 import numpy as np
+from norch.autograd.functions import *
 
 def sigmoid(x):
-    return 1.0 / (1.0 + (math.e) ** (-x)) 
+    z = x.sigmoid()
+    return z
 
 def softmax(x, dim=None):
     if dim is not None and dim < 0:

norch/tensor.py

@@ -948,6 +948,25 @@ class Tensor:
         
         return result_data
     
+    def sigmoid(self):
+        Tensor._C.sigmoid_tensor.argtypes = [ctypes.POINTER(CTensor)]
+        Tensor._C.sigmoid_tensor.restype = ctypes.POINTER(CTensor)
+
+        result_tensor_ptr = Tensor._C.sigmoid_tensor(self.tensor)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = self.shape.copy()
+        result_data.ndim = self.ndim
+        result_data.device = self.device
+        result_data.numel = self.numel
+
+        result_data.requires_grad = self.requires_grad
+        if result_data.requires_grad:
+            result_data.grad_fn = SigmoidBackward(self)
+        
+        return result_data
+    
 
     
     def transpose(self, axis1, axis2):

test.py

@@ -81,9 +81,9 @@ train_loader = Dataloader(train_data, batch_size = BATCH_SIZE)
 class MyModel(nn.Module):
     def __init__(self):
         super(MyModel, self).__init__()
-        self.fc1 = nn.Linear(784, 5)
+        self.fc1 = nn.Linear(784, 10)
         self.sigmoid1 = nn.Sigmoid()
-        self.fc2 = nn.Linear(5, 10)
+        self.fc2 = nn.Linear(10, 10)
         self.sigmoid2 = nn.Sigmoid()
 
     def forward(self, x):

tests/test_autograd.py

@@ -601,7 +601,7 @@ class TestTensorAutograd(unittest.TestCase):
 
         torch_result.backward()
         torch_tensor_grad = torch_tensor.grad
-        
+
         self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
 
     def test_mse_loss_autograd(self):

tests/test_operations.py

@@ -258,10 +258,6 @@ class TestTensorOperations(unittest.TestCase):
         torch_expected_0 = torch_tensor.squeeze(0)
         self.assertTrue(utils.compare_torch(torch_squeeze_0, torch_expected_0))
 
-        # Squeeze at dim=2 (should raise an error because size is not 1)
-        with self.assertRaises(ValueError):
-            norch_tensor.squeeze(2)
-
         # Create a tensor with a dimension of size 1 in the middle
         norch_tensor_middle_1 = norch.Tensor([[[1, 2]], [[3, 4]]]).to(self.device)  # shape [2, 1, 2]