Implement iterable datasets support. (#1095)

* Implement iterable datasets support.

* Don't allow `NULLs` to be propagated in the dataloader.

NAMESPACE

@@ -69,6 +69,7 @@ S3method(as.matrix,torch_tensor)
 S3method(as.numeric,torch_tensor)
 S3method(as_array,torch_tensor)
 S3method(as_iterator,dataloader)
+S3method(as_iterator,iterable_dataset)
 S3method(as_iterator,utils_sampler)
 S3method(asin,torch_tensor)
 S3method(atan,torch_tensor)
@@ -83,6 +84,7 @@ S3method(expm1,torch_tensor)
 S3method(floor,torch_tensor)
 S3method(length,dataloader)
 S3method(length,dataset)
+S3method(length,iterable_dataset)
 S3method(length,nn_module_list)
 S3method(length,nn_sequential)
 S3method(length,torch_tensor)
@@ -100,6 +102,7 @@ S3method(nn_prune_head,nn_sequential)
 S3method(print,R7)
 S3method(print,cuda_memory_stats)
 S3method(print,dataset_generator)
+S3method(print,iterable_dataset_generator)
 S3method(print,nn_module)
 S3method(print,script_function)
 S3method(print,script_method)
@@ -178,6 +181,7 @@ export(is_torch_layout)
 export(is_torch_memory_format)
 export(is_torch_qscheme)
 export(is_undefined_tensor)
+export(iterable_dataset)
 export(jit_compile)
 export(jit_load)
 export(jit_ops)
@@ -526,9 +530,11 @@ export(torch_ceil)
 export(torch_celu)
 export(torch_celu_)
 export(torch_cfloat)
+export(torch_cfloat128)
 export(torch_cfloat32)
 export(torch_cfloat64)
 export(torch_chain_matmul)
+export(torch_chalf)
 export(torch_channel_shuffle)
 export(torch_channels_last_format)
 export(torch_cholesky)

R/utils-data-dataloader.R

@@ -152,11 +152,15 @@ DataLoader <- R6::R6Class(
 
       if (is_map_dataset(dataset)) {
         self$.dataset_kind <- "map"
+      } else if (is_iterable_dataset(dataset)) {
+        self$.dataset_kind <- "iterable"
+      } else {
+        cli::cli_abort("Unknown dataset type with class {.cls {class(dataset)}}")
       }
 
       if (is.null(sampler)) {
         if (self$.dataset_kind == "iterable") {
-          # TODO
+          sampler <- InfiniteSampler()
         } else {
           if (shuffle) {
             sampler <- RandomSampler(dataset, generator = generator)
@@ -200,13 +204,26 @@ DataLoader <- R6::R6Class(
         }
 
         MultiProcessingDataLoaderIter$new(self)
+      } else if (self$.dataset_kind == "iterable") {
+        if (self$num_workers == 0) {
+          return(SingleProcessDataLoaderIter$new(self))
+        }
+        cli::cli_abort("Multi-process dataloader not implemented yet for Iterable datasets.")
       } else {
         not_implemented_error()
       }
     },
     .length = function() {
       if (self$.dataset_kind == "iterable") {
-        not_implemented_error()
+        l <- length(self$dataset)
+
+        if (is.na(l)) return(l)
+
+        if (self$drop_last) {
+          return(l %/% self$batch_size)
+        } else {
+          return(as.integer(ceiling(l / self$batch_size)))
+        } 
       } else {
         length(self$.index_sampler)
       }
@@ -283,6 +300,13 @@ SingleProcessDataLoaderIter <- R6::R6Class(
           self$.collate_fn,
           self$.drop_last
         )
+      } else if (self$.dataset_kind == "iterable") {
+        self$.dataset_fetcher <- IterableDatasetFetcher$new(
+          self$.dataset,
+          self$.auto_collation,
+          self$.collate_fn,
+          self$.drop_last
+        )
       } else {
         not_implemented_error()
       }
@@ -294,7 +318,9 @@ SingleProcessDataLoaderIter <- R6::R6Class(
         return(coro::exhausted())
       }
 
+      # data can be exhausted in iterable datasets
       data <- self$.dataset_fetcher$fetch(index)
+
       if (self$.pin_memory) {
         # TODO
       }

R/utils-data-fetcher.R

@@ -36,6 +36,12 @@ IterableDatasetFetcher <- R6::R6Class(
           d <- self$dataset_iter()
 
           if (is_exhausted(d)) {
+            if (self$drop_last || i ==  1) {
+              return(coro::exhausted())
+            }
+
+            # we drop the null values in that list.
+            data <- data[seq_len(i-1L)]
             break
           }
 
@@ -44,6 +50,7 @@ IterableDatasetFetcher <- R6::R6Class(
       } else {
         data <- self$dataset_iter()
       }
+
       self$collate_fn(data)
     }
   )

R/utils-data-sampler.R

@@ -159,6 +159,19 @@ BatchSampler <- sampler(
   }
 )
 
+InfiniteSampler <- sampler(
+  "infinite_sampler", 
+  initialize = function() {},
+  .iter = function() {
+    function() {
+      TRUE
+    }
+  },
+  .length = function() {
+    Inf
+  }
+)
+
 #' @export
 as_iterator.utils_sampler <- function(x) {
   it <- x$.iter()

R/utils-data.R

@@ -36,10 +36,32 @@ Dataset <- R6::R6Class(
   )
 )
 
+IterableDataset <- R6::R6Class(
+  classname = "iterable_dataset",
+  lock_objects = FALSE,
+  public = list(
+    .iter = function() {
+      not_implemented_error()
+    },
+    .length = function() {
+      NA_integer_
+    }
+  )
+)
+
 is_map_dataset <- function(x) {
   inherits(x, "dataset")
 }
 
+is_iterable_dataset <- function(x) {
+  inherits(x, "iterable_dataset")
+}
+
+#' @export
+as_iterator.iterable_dataset <- function(x) {
+  x$.iter()
+}
+
 get_init <- function(x) {
   if (!is.null(x$public_methods$initialize)) {
     return(x$public_methods$initialize)
@@ -107,12 +129,58 @@ dataset <- function(name = NULL, inherit = Dataset, ...,
   )
 }
 
+
+#' Creates an iterable dataset
+#' 
+#' @inheritParams dataset
+#' @examples
+#' ids <- iterable_dataset(
+#'   name = "hello",
+#'   initialize = function(n = 5) {
+#'     self$n <- n
+#'     self$i <- 0
+#'   },
+#'   .iter = function() {
+#'     i <- 0
+#'     function() {
+#'       i <<- i + 1
+#'       if (i > self$n) {
+#'         coro::exhausted()
+#'       } else {
+#'         i
+#'       }
+#'     }
+#'   }
+#' )
+#' coro::collect(ids()$.iter())
+#' @export
+iterable_dataset <- function(name, inherit = IterableDataset, ..., 
+                             private = NULL, active = NULL,
+                             parent_env = parent.frame()) {
+  create_class(
+    name = name,
+    inherit = inherit,
+    ...,
+    private = private,
+    active = active,
+    parent_env = parent_env,
+    attr_name = "Dataset",
+    constructor_class = "iterable_dataset_generator"
+  )
+}
+
 #' @export
 print.dataset_generator <- function(x, ...) {
   cli::cat_line("<dataset_generator>")
   print(attr(x, "Dataset"))
 }
 
+#' @export
+print.iterable_dataset_generator <- function(x, ...) {
+  cli::cat_line("<iterable_dataset_generator>")
+  print(attr(x, "IterableDataset"))
+}
+
 #' @export
 `[.dataset` <- function(x, y) {
   y <- as.integer(y)
@@ -136,6 +204,11 @@ length.dataset <- function(x) {
   x$.length()
 }
 
+#' @export
+length.iterable_dataset <- function(x) {
+  x$.length()
+}
+
 #' Dataset wrapping tensors.
 #'
 #' Each sample will be retrieved by indexing tensors along the first dimension.

src/utils.cpp

@@ -65,16 +65,19 @@ std::thread::id main_thread_id() noexcept {
 // [[Rcpp::export]]
 Rcpp::List transpose2(Rcpp::List x) {
   auto templ = Rcpp::as<Rcpp::List>(x[0]);
-  auto num_elements = templ.length();
+  const auto num_elements = templ.length();
 
-  auto size = x.length();
+  const auto size = x.length();
   std::vector<Rcpp::List> out;
 
   for (auto i = 0; i < num_elements; i++) {
     out.push_back(Rcpp::List(size));
   }
 
   for (size_t j = 0; j < size; j++) {
+    if (Rf_isNull(x[j])) {
+      Rcpp::stop("NULL is not allowed. Expected a list.");
+    }
     auto el = Rcpp::as<Rcpp::List>(x[j]);
     for (auto i = 0; i < num_elements; i++) {
       out[i][j] = el[i];