Repository: lhnguyen102/cuTAGI (language: python)

Issue:
Problem with the MixtureRelu, MixtureTanh, MixtureSigmoid and CUDA
I started implementing the new formulation for the mixture-based activation functions and I realized that there an issue with the `MixtureRelu`, `MixtureTanh`, `MixtureSigmoid` while running pytagi_v1 on CUDA. 

If I run the [`test.py`](pytagi/pytagi_v1/test.py) example on CPU, all the activation function work properly. When running with CUDA the MixtureRelu, MixtureTanh and MixtureSigmoid do not learn. If I run the same test on CUDA with the standard ReLU it works properly.

If I run the c++ example [`cfg_mnist_2fc.txt`](cfg/cfg_mnist_2fc.txt), all mixture-based activations work well on CUDA. 

@lhnguyen102 Any insights about what could cause this issue?
Problem with the MixtureRelu, MixtureTanh, MixtureSigmoid and CUDA
I started implementing the new formulation for the mixture-based activation functions and I realized that there an issue with the `MixtureRelu`, `MixtureTanh`, `MixtureSigmoid` while running pytagi_v1 on CUDA. 

If I run the [`test.py`](pytagi/pytagi_v1/test.py) example on CPU, all the activation function work properly. When running with CUDA the MixtureRelu, MixtureTanh and MixtureSigmoid do not learn. If I run the same test on CUDA with the standard ReLU it works properly.

If I run the c++ example [`cfg_mnist_2fc.txt`](cfg/cfg_mnist_2fc.txt), all mixture-based activations work well on CUDA. 

@lhnguyen102 Any insights about what could cause this issue?


Discussion / hints:

Fix patch:
diff --git a/src/activation_cuda.cu b/src/activation_cuda.cu
index fc1409e09..dfbfb7512 100644
--- a/src/activation_cuda.cu
+++ b/src/activation_cuda.cu
@@ -3,7 +3,7 @@
 // Description:  ...
 // Authors:      Luong-Ha Nguyen & James-A. Goulet
 // Created:      December 04, 2023
-// Updated:      March 31, 2024
+// Updated:      April 02, 2024
 // Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
 // License:      This code is released under the MIT License.
 ////////////////////////////////////////////////////////////////////////////////
@@ -196,8 +196,6 @@ void TanhCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     int num_states = input_states.actual_size * input_states.block_size;
     unsigned int blocks =
         (num_states + this->num_cuda_threads - 1) / this->num_cuda_threads;
@@ -269,8 +267,6 @@ void MixtureReluCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     int num_states = input_states.actual_size * input_states.block_size;
     unsigned int blocks =
         (num_states + this->num_cuda_threads - 1) / this->num_cuda_threads;
@@ -342,8 +338,6 @@ void MixtureSigmoidCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     int num_states = input_states.actual_size * input_states.block_size;
     unsigned int blocks =
         (num_states + this->num_cuda_threads - 1) / this->num_cuda_threads;
@@ -415,8 +409,6 @@ void MixtureTanhCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     int num_states = input_states.actual_size * input_states.block_size;
     unsigned int blocks =
         (num_states + this->num_cuda_threads - 1) / this->num_cuda_threads;
@@ -488,8 +480,6 @@ void SoftplusCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     int num_states = input_states.actual_size * input_states.block_size;
     unsigned int blocks =
         (num_states + this->num_cuda_threads - 1) / this->num_cuda_threads;
@@ -634,8 +624,6 @@ void SoftmaxCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
-    cu_input_states->to_device();
-
     unsigned int blocks =
         (input_states.block_size + this->num_cuda_threads - 1) /
         this->num_cuda_threads;


Test patch:
diff --git a/test/fnn/test_fnn_mnist_cpu.cpp b/test/fnn/test_fnn_mnist_cpu.cpp
index 9a54b50a8..df5513b5c 100644
--- a/test/fnn/test_fnn_mnist_cpu.cpp
+++ b/test/fnn/test_fnn_mnist_cpu.cpp
@@ -66,8 +66,8 @@ void fnn_mnist() {
     //////////////////////////////////////////////////////////////////////
     // TAGI network
     //////////////////////////////////////////////////////////////////////
-    Sequential model(Linear(784, 100), ReLU(), Linear(100, 100), ReLU(),
-                     Linear(100, 11));
+    Sequential model(Linear(784, 100), MixtureRelu(), Linear(100, 100),
+                     MixtureRelu(), Linear(100, 11));
 
     // Sequential model(Linear(784, 100), BatchNorm2d(100), ReLU(),
     //                  Linear(100, 100), BatchNorm2d(100), ReLU(),
@@ -97,7 +97,7 @@ void fnn_mnist() {
     //                  Linear(100, 11));
 
     // model.set_threads(8);
-    // model.to_device("cuda");
+    model.to_device("cuda");
     // model.preinit_layer();
     // model.load("test_model/test_model.bin");
 
@@ -173,7 +173,7 @@ void fnn_mnist() {
         std::cout << "Epoch #" << e + 1 << "/" << n_epochs << "\n";
         std::cout << "Training...\n";
         auto start = std::chrono::steady_clock::now();
-        for (int i = 0; i < 1; i++) {
+        for (int i = 0; i < iters; i++) {
             // Load data
             get_batch_images_labels(train_db, data_idx, batch_size, i, x_batch,
                                     y_batch, idx_ud_batch, label_batch);