fix support for MPS Apple

Author: gilbertocamara

DESCRIPTION

@@ -246,7 +246,6 @@ Collate:
     'sits_reclassify.R'
     'sits_reduce.R'
     'sits_regularize.R'
-    'sits_resnet.R'
     'sits_sample_functions.R'
     'sits_segmentation.R'
     'sits_select.R'

R/api_torch.R

@@ -372,8 +372,7 @@
     torch::cuda_is_available()
 }
 .torch_has_mps <- function(){
-    # torch::backends_mps_is_available()
-    return(FALSE)
+    torch::backends_mps_is_available()
 }
 
 .torch_mem_info <- function() {
@@ -400,8 +399,9 @@
 #' @return TRUE/FALSE
 #'
 .torch_gpu_enabled <- function(ml_model){
-    gpu_enabled <- (inherits(ml_model, "torch_model") &&
-        (.torch_has_cuda() || .torch_has_mps())
+    gpu_enabled <- (
+        inherits(ml_model, "torch_model") &&
+        .torch_has_cuda()
     )
     return(gpu_enabled)
 }
@@ -439,7 +439,7 @@
     return(values)
 }
 #' @title Use GPU or CPU train for MPS Apple
-#' @name .torch_mps_train
+#' @name .torch_cpu_train
 #' @author Gilberto Camara, \email{gilberto.camara@@inpe.br}
 #' @keywords internal
 #' @noRd
@@ -448,12 +448,11 @@
 #'
 #' @return TRUE/FALSE
 #'
-.torch_mps_train <- function() {
-    if (torch::backends_mps_is_available())
-        cpu_train <-  TRUE
-    else
+.torch_cpu_train <- function() {
+    if (torch::cuda_is_available())
         cpu_train <-  FALSE
-
+    else
+        cpu_train <-  TRUE
     return(cpu_train)
 }
 

R/sits_classify.R

@@ -233,11 +233,18 @@ sits_classify.raster_cube <- function(data,
         )
         # Calculate available memory from GPU
         memsize <- floor(gpu_memory - .torch_mem_info())
-        .check_int_parameter(memsize, min = 2,
+        .check_int_parameter(memsize, min = 1,
                         msg = .conf("messages", ".check_gpu_memory_size")
         )
         proc_bloat <- .conf("processing_bloat_gpu")
     }
+    # avoid memory race in Apple MPS
+    if(.torch_has_mps()){
+        memsize <- 1
+        gpu_memory <- 1
+    }
+    # save memsize for latter use
+    sits_env[["gpu_memory"]] <- gpu_memory
     # Spatial filter
     if (.has(roi)) {
         roi <- .roi_as_sf(roi)

R/sits_lighttae.R

@@ -271,7 +271,7 @@ sits_lighttae <- function(samples = NULL,
             }
         )
         # torch 12.0 not working with Apple MPS
-        cpu_train <- .torch_mps_train()
+        cpu_train <- .torch_cpu_train()
         # Train the model using luz
         torch_model <-
             luz::setup(
@@ -339,17 +339,26 @@ sits_lighttae <- function(samples = NULL,
             values <- array(
                 data = as.matrix(values), dim = c(n_samples, n_times, n_bands)
             )
-            # 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")
-            )
-            # Are the results consistent with the data input?
-            .check_processed_values(
-                values = values, input_pixels = input_pixels
-            )
+            # Load into GPU
+            if (.torch_has_cuda()){
+                # set the batch size according to the GPU memory
+                gpu_memory <- sits_env[["gpu_memory"]]
+                b_size <- 2^gpu_memory
+                # transfor the input array to a dataset
+                values <- .as_dataset(values)
+                # To the data set to a torcj  transform in a dataloader to use the batch size
+                values <- torch::dataloader(values, batch_size = b_size)
+                # Do GPU classification with dataloader
+                values <- .try(
+                    stats::predict(object = torch_model, values),
+                    .msg_error = .conf("messages", ".check_gpu_memory_size")
+                )
+            } else {
+                # Do  classification without dataloader
+                values <- stats::predict(object = torch_model, values)
+            }
+            # Convert from tensor to array
+            values <- torch::as_array(values)
             # Update the columns names to labels
             colnames(values) <- labels
             return(values)

R/sits_mlp.R

@@ -205,7 +205,8 @@ sits_mlp <- function(samples = NULL,
         test_y <- unname(code_labels[.pred_references(test_samples)])
         # Set torch seed
         torch::torch_manual_seed(sample.int(10^5, 1))
-
+        # train with CPU or GPU?
+        cpu_train <- .torch_cpu_train()
         # Define the MLP architecture
         mlp_model <- torch::nn_module(
             initialize = function(num_pred, layers, dropout_rates, y_dim) {
@@ -241,8 +242,8 @@ sits_mlp <- function(samples = NULL,
                 self$model(x)
             }
         )
-        # torch 12.0 not working with Apple MPS
-        cpu_train <- .torch_mps_train()
+        # Train with CPU or GPU?
+        cpu_train <- .torch_cpu_train()
         # Train the model using luz
         torch_model <-
             luz::setup(
@@ -290,16 +291,27 @@ sits_mlp <- function(samples = NULL,
             values <- .pred_normalize(pred = values, stats = ml_stats)
             # Transform input into matrix
             values <- as.matrix(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")
-            )
-            # Are the results consistent with the data input?
-            .check_processed_values(
-                values = values, input_pixels = input_pixels
-            )
+            # if CUDA is available, transform to torch data set
+            # Load into GPU
+            if (.torch_has_cuda()){
+                # set the batch size according to the GPU memory
+                gpu_memory <- sits_env[["gpu_memory"]]
+                b_size <- 2^gpu_memory
+                # transfor the input array to a dataset
+                values <- .as_dataset(values)
+                # To the data set to a torcj  transform in a dataloader to use the batch size
+                values <- torch::dataloader(values, batch_size = b_size)
+                # Do GPU classification with dataloader
+                values <- .try(
+                    stats::predict(object = torch_model, values),
+                    .msg_error = .conf("messages", ".check_gpu_memory_size")
+                )
+            } else {
+                # Do  classification without dataloader
+                values <- stats::predict(object = torch_model, values)
+            }
+            # Convert from tensor to array
+            values <- torch::as_array(values)
             # Update the columns names to labels
             colnames(values) <- labels
             return(values)