improve api_torch and disable MPS

Author: gilbertocamara

R/api_classify.R

@@ -195,10 +195,8 @@
         start_time = tile_start_time,
         verbose = verbose
     )
-    # Clean torch allocations
-    if (.is_torch_model(ml_model) && .torch_has_cuda()) {
-        torch::cuda_empty_cache()
-    }
+    # Clean GPU memory allocation
+    .ml_gpu_clean(ml_model)
     # Return probs tile
     probs_tile
 }
@@ -527,9 +525,7 @@
         )
     }
     # choose between GPU and CPU
-    if (inherits(ml_model, "torch_model") &&
-        (.torch_has_cuda() || .torch_has_mps())
-    )
+    if (.torch_gpu_enabled(ml_model))
         prediction <- .classify_ts_gpu(
             pred = pred,
             ml_model = ml_model,
@@ -630,10 +626,8 @@
         values <- ml_model(values)
         # Return classification
         values <- tibble::tibble(data.frame(values))
-        # Clean torch cache
-        if (.is_torch_model(ml_model) && .torch_has_cuda()) {
-            torch::cuda_empty_cache()
-        }
+        # Clean GPU memory
+        .ml_gpu_clean(ml_model)
         return(values)
     })
     return(prediction)

R/api_ml_model.R

@@ -93,3 +93,15 @@
     names(labels) <- seq_along(labels)
     labels
 }
+#' @title Clean GPU memory allocation
+#' @keywords internal
+#' @noRd
+#' @param  ml_model  Closure that contains ML model and its environment
+#' @return           Called for side effects
+.ml_gpu_clean <- function(ml_model) {
+    # Clean torch allocations
+    if (.is_torch_model(ml_model) && .torch_has_cuda()) {
+        torch::cuda_empty_cache()
+    }
+    return(invisible(NULL))
+}

R/api_torch.R

@@ -372,7 +372,8 @@
     torch::cuda_is_available()
 }
 .torch_has_mps <- function(){
-    torch::backends_mps_is_available()
+    # torch::backends_mps_is_available()
+    return(FALSE)
 }
 
 .torch_mem_info <- function() {
@@ -386,6 +387,75 @@
     }
     return(mem_sum)
 }
+#' @title Verify if torch works on GPU
+#' @name .torch_gpu_enabled
+#' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
+#' @keywords internal
+#' @noRd
+#' @description Use CPU or GPU for torch models depending on
+#' availability
+#'
+#' @param ml_model   ML model
+#'
+#' @return TRUE/FALSE
+#'
+.torch_gpu_enabled <- function(ml_model){
+    gpu_enabled <- (inherits(ml_model, "torch_model") &&
+        (.torch_has_cuda() || .torch_has_mps())
+    )
+    return(gpu_enabled)
+}
+#' @title Torch function for prediction in GPU or CPU
+#' @name .torch_predict
+#' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
+#' @keywords internal
+#' @noRd
+#' @description Use CPU or GPU for torch models depending on
+#' availability
+#'
+#' @param values        Values to be predicted
+#' @param torch_model   Torch model
+#'
+#' @return Predicted values
+#'
+.torch_predict <- function(values, torch_model){
+    # if CUDA is available, transform to torch data set
+    # Load into GPU
+    if (.torch_has_cuda() || .torch_has_mps()) {
+        values <- .as_dataset(values)
+        # We need to transform in a dataloader to use the batch size
+        values <- torch::dataloader(
+            values, batch_size = 2^15
+        )
+        # Do GPU classification
+        values <- .try(
+            stats::predict(object = torch_model, values),
+            .msg_error = .conf("messages", ".check_gpu_memory_size")
+        )
+    } else {
+        # Do CPU classification
+        values <- stats::predict(object = torch_model, values)
+    }
+    return(values)
+}
+#' @title Use GPU or CPU train for MPS Apple
+#' @name .torch_mps_train
+#' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
+#' @keywords internal
+#' @noRd
+#' @description Use CPU or GPU for torch models depending on
+#' availability
+#'
+#' @return TRUE/FALSE
+#'
+.torch_mps_train <- function() {
+    if (torch::backends_mps_is_available())
+        cpu_train <-  TRUE
+    else
+        cpu_train <-  FALSE
+
+    return(cpu_train)
+}
 
 .as_dataset <- torch::dataset(
     "dataset",

R/sits_classify.R

@@ -228,8 +228,7 @@ sits_classify.raster_cube <- function(data,
     # Get default proc bloat
     proc_bloat <- .conf("processing_bloat_cpu")
     # If we using the GPU, gpu_memory parameter needs to be specified
-    if (.is_torch_model(ml_model) &&
-        (.torch_has_cuda() || .torch_has_mps())) {
+    if (.torch_gpu_enabled(ml_model)) {
         .check_int_parameter(gpu_memory, min = 1, max = 16384,
                              msg = .conf("messages", ".check_gpu_memory")
         )
@@ -356,7 +355,7 @@ sits_classify.segs_cube <- function(data,
                                     end_date = NULL,
                                     memsize = 8L,
                                     multicores = 2L,
-                                    gpu_memory = 16,
+                                    gpu_memory = 4,
                                     output_dir,
                                     version = "v1",
                                     n_sam_pol = NULL,
@@ -377,8 +376,7 @@ sits_classify.segs_cube <- function(data,
     .check_progress(progress)
     proc_bloat <- .conf("processing_bloat_seg_class")
     # If we using the GPU, gpu_memory parameter needs to be specified
-    if (.is_torch_model(ml_model) &&
-        (.torch_has_cuda() || .torch_has_mps())) {
+    if (.torch_gpu_enabled(ml_model)) {
         .check_int_parameter(gpu_memory, min = 1, max = 16384,
                              msg = .conf("messages", ".check_gpu_memory")
         )

R/sits_lighttae.R

@@ -271,10 +271,7 @@ sits_lighttae <- function(samples = NULL,
             }
         )
         # torch 12.0 not working with Apple MPS
-        if (torch::backends_mps_is_available())
-            cpu_train <-  TRUE
-        else
-            cpu_train <-  FALSE
+        cpu_train <- .torch_mps_train()
         # Train the model using luz
         torch_model <-
             luz::setup(
@@ -342,23 +339,9 @@ sits_lighttae <- function(samples = NULL,
             values <- array(
                 data = as.matrix(values), dim = c(n_samples, n_times, n_bands)
             )
-            # if CUDA is available, transform to torch data set
-            # Load into GPU
-            if (.torch_has_cuda() || .torch_has_mps()) {
-                values <- .as_dataset(values)
-                # We need to transform in a dataloader to use the batch size
-                values <- torch::dataloader(
-                    values, batch_size = 2^15
-                )
-                # Do GPU classification
-                values <- .try(
-                    stats::predict(object = torch_model, values),
-                    .msg_error = .conf("messages", ".check_gpu_memory_size")
-                )
-            } else {
-                # Do CPU classification
-                values <- stats::predict(object = torch_model, values)
-            }
+            # Predict using GPU if available
+            # If not, use CPU
+            values <- .torch_predict(values, torch_model)
             # Convert to tensor CPU
             values <- torch::as_array(
                 x = torch::torch_tensor(values, device = "cpu")

R/sits_mlp.R

@@ -58,7 +58,8 @@
 #' @examples
 #' if (sits_run_examples()) {
 #'     # create an MLP model
-#'     torch_model <- sits_train(samples_modis_ndvi, sits_mlp())
+#'     torch_model <- sits_train(samples_modis_ndvi,
+#'            sits_mlp(epochs = 20, verbose = TRUE))
 #'     # plot the model
 #'     plot(torch_model)
 #'     # create a data cube from local files
@@ -241,10 +242,7 @@ sits_mlp <- function(samples = NULL,
             }
         )
         # torch 12.0 not working with Apple MPS
-        if (torch::backends_mps_is_available())
-            cpu_train <-  TRUE
-        else
-            cpu_train <-  FALSE
+        cpu_train <- .torch_mps_train()
         # Train the model using luz
         torch_model <-
             luz::setup(
@@ -292,23 +290,8 @@ sits_mlp <- function(samples = NULL,
             values <- .pred_normalize(pred = values, stats = ml_stats)
             # Transform input into matrix
             values <- as.matrix(values)
-            # if CUDA is available, transform to torch data set
-            # Load into GPU
-            if (.torch_has_cuda() || .torch_has_mps()) {
-                values <- .as_dataset(values)
-                # We need to transform in a dataloader to use the batch size
-                values <- torch::dataloader(
-                    values, batch_size = 2^15
-                )
-                # Do GPU classification
-                values <- .try(
-                    stats::predict(object = torch_model, values),
-                    .msg_error = .conf("messages", ".check_gpu_memory_size")
-                )
-            } else {
-                # Do CPU classification
-                values <- stats::predict(object = torch_model, values)
-            }
+            # predict using CPU or GPU depending on machine
+            values <- .torch_predict(values, torch_model)
             # Convert to tensor cpu to support GPU processing
             values <- torch::as_array(
                 x = torch::torch_tensor(values, device = "cpu")

R/sits_tae.R

@@ -244,10 +244,7 @@ sits_tae <- function(samples = NULL,
             }
         )
         # torch 12.0 not working with Apple MPS
-        if (.torch_has_mps())
-            cpu_train <-  TRUE
-        else
-            cpu_train <-  FALSE
+        cpu_train <- .torch_mps_train()
         # train the model using luz
         torch_model <-
             luz::setup(
@@ -309,23 +306,9 @@ sits_tae <- function(samples = NULL,
             values <- array(
                 data = as.matrix(values), dim = c(n_samples, n_times, n_bands)
             )
-            # if CUDA is available, transform to torch data set
-            # Load into GPU
-            if (.torch_has_cuda() || .torch_has_mps()) {
-                values <- .as_dataset(values)
-                # We need to transform in a dataloader to use the batch size
-                values <- torch::dataloader(
-                    values, batch_size = 2^15
-                )
-                # Do GPU classification
-                values <- .try(
-                    stats::predict(object = torch_model, values),
-                    .msg_error = .conf("messages", ".check_gpu_memory_size")
-                )
-            } else {
-                # Do CPU classification
-                values <- stats::predict(object = torch_model, values)
-            }
+            # Predict using GPU if available
+            # If not, use CPU
+            values <- .torch_predict(values, torch_model)
             # Convert to tensor CPU
             values <- torch::as_array(
                 x = torch::torch_tensor(values, device = "cpu")

R/sits_tempcnn.R

@@ -60,7 +60,7 @@
 #' if (sits_run_examples()) {
 #'     # create a TempCNN model
 #'     torch_model <- sits_train(samples_modis_ndvi,
-#'                sits_tempcnn(verbose = TRUE))
+#'                sits_tempcnn(epochs = 20, verbose = TRUE))
 #'     # plot the model
 #'     plot(torch_model)
 #'     # create a data cube from local files
@@ -290,10 +290,7 @@ sits_tempcnn <- function(samples = NULL,
             }
         )
         # torch 12.0 not working with Apple MPS
-        if (torch::backends_mps_is_available())
-            cpu_train <-  TRUE
-        else
-            cpu_train <-  FALSE
+        cpu_train <- .torch_mps_train()
         # Train the model using luz
         torch_model <-
             luz::setup(
@@ -362,22 +359,8 @@ sits_tempcnn <- function(samples = NULL,
                 data = as.matrix(values), dim = c(n_samples, n_times, n_bands)
             )
             # if CUDA is available, transform to torch data set
-            # Load into GPU
-            if (.torch_has_cuda() || .torch_has_mps()) {
-                values <- .as_dataset(values)
-                # We need to transform in a dataloader to use the batch size
-                 values <- torch::dataloader(
-                    values, batch_size = 2^15
-                )
-                # Do GPU classification
-                values <- .try(
-                    stats::predict(object = torch_model, values),
-                    .msg_error = .conf("messages", ".check_gpu_memory_size")
-                )
-            } else {
-                # Do CPU classification
-                values <- stats::predict(object = torch_model, values)
-            }
+            # predict using CPU or GPU depending on machine
+            values <- .torch_predict(values, torch_model)
             # Convert to tensor cpu to support GPU processing
             values <- torch::as_array(
                 x = torch::torch_tensor(values, device = "cpu")