feat/pipeop-nn-ft-cls (#381)

Author: cxzhang4

DESCRIPTION

@@ -117,6 +117,7 @@ Collate:
     'PipeOpTorchConv.R'
     'PipeOpTorchConvTranspose.R'
     'PipeOpTorchDropout.R'
+    'PipeOpTorchFTCLS.R'
     'PipeOpTorchFn.R'
     'PipeOpTorchHead.R'
     'PipeOpTorchIdentity.R'

NAMESPACE

@@ -102,6 +102,7 @@ export(PipeOpTorchConvTranspose2D)
 export(PipeOpTorchConvTranspose3D)
 export(PipeOpTorchDropout)
 export(PipeOpTorchELU)
+export(PipeOpTorchFTCLS)
 export(PipeOpTorchFlatten)
 export(PipeOpTorchFn)
 export(PipeOpTorchGELU)
@@ -181,6 +182,7 @@ export(model_descriptor_to_learner)
 export(model_descriptor_to_module)
 export(model_descriptor_union)
 export(nn)
+export(nn_ft_cls)
 export(nn_geglu)
 export(nn_graph)
 export(nn_merge_cat)

NEWS.md

@@ -5,6 +5,7 @@
 * feat: `TorchIngressToken` now also can take a `Selector` as argument `features`.
 * feat: Added encoders for numericals and categoricals
 * feat: Added `po("nn_fn")` for calling custom functions in a network.
+* feat: Added `po("nn_ft_cls")` for concatenating a CLS token to a tokenized input.
 
 # mlr3torch 0.2.1
 

R/PipeOpTorchFTCLS.R

@@ -0,0 +1,88 @@
+#' @title CLS Token for FT-Transformer
+#' @description
+#' Concatenates a CLS token to the input as the last feature.
+#' The input shape is expected to be `(batch, n_features, d_token)` and the output shape is
+#' `(batch, n_features + 1, d_token)`.
+#'
+#' This is used in the FT-Transformer.
+#'
+#' @param d_token (`integer(1)`)\cr
+#'   The dimension of the embedding.
+#' @param initialization (`character(1)`)\cr
+#'   The initialization method for the embedding weights. Possible values are `"uniform"`
+#'   and `"normal"`.
+#'
+#' @references
+#' `r format_bib("devlin2018bert")`
+#'
+#' @export
+nn_ft_cls = nn_module(
+  "nn_ft_cls",
+  initialize = function(d_token, initialization) {
+    self$d_token = d_token
+    # an individual CLS token
+    self$weight = nn_parameter(torch_empty(d_token))
+    self$initialization = initialization
+    self$reset_parameters()
+  },
+  reset_parameters = function() {
+    initialize_token_(self$weight, d = self$d_token, self$initialization)
+  },
+  # Repeats the underlying CLS token to create a tensor with the given leading dimensions.
+  # Used for creating a batch of CLS tokens
+  expand = function(...) {
+    leading_dimensions = list(...)
+    if (length(leading_dimensions) == 0) {
+      return(self$weight)
+    }
+    new_dims = rep(1, length(leading_dimensions) - 1)
+    return(self$weight$view(c(new_dims, -1))$expand(c(leading_dimensions, -1)))
+  },
+  forward = function(input) {
+    return(torch_cat(list(input, self$expand(input$shape[1], 1)), dim = 2))
+  }
+)
+
+#' @title CLS Token for FT-Transformer
+#' @inherit nn_ft_cls description
+#' @section nn_module:
+#' Calls [`nn_ft_cls()`] when trained.
+#' @templateVar id nn_ft_cls
+#' @template pipeop_torch
+#' @template pipeop_torch_example
+#' @export
+PipeOpTorchFTCLS = R6::R6Class("PipeOpTorchFTCLS",
+inherit = PipeOpTorch,
+  public = list(
+    #' @description Creates a new instance of this [R6][R6::R6Class] class.
+    #' @template params_pipelines
+    initialize = function(id = "nn_ft_cls", param_vals = list()) {
+      param_set = ps(
+        initialization = p_fct(tags = c("train"), levels = c("uniform", "normal"), default = "uniform")
+      )
+
+      super$initialize(
+        id = id,
+        module_generator = nn_ft_cls,
+        param_vals = param_vals,
+        param_set = param_set
+      )
+    }
+  ),
+  private = list(
+    .shapes_out = function(shapes_in, param_vals, task) {
+      if (length(shapes_in$input) != 3) {
+        stop("Input tensor must have 3 dimensions.")
+      }
+      shapes_in[[1]][2] = shapes_in[[1]][2] + 1
+      return(shapes_in)
+    },
+    .shape_dependent_params = function(shapes_in, param_vals, task) {
+      param_vals$d_token = shapes_in$input[3]
+      return(param_vals)
+    }
+  )
+)
+
+#' @include aaa.R
+register_po("nn_ft_cls", PipeOpTorchFTCLS)

R/PipeOpTorchTokenizer.R

@@ -68,6 +68,7 @@ PipeOpTorchTokenizerNum = R6Class("PipeOpTorchTokenizerNum",
 #' @name nn_tokenizer_num
 #' @description
 #' Tokenizes numeric features into a dense embedding.
+#' For an input of shape `(batch, n_features)` the output shape is `(batch, n_features, d_token)`.
 #' @param n_features (`integer(1)`)\cr
 #'   The number of features.
 #' @param d_token (`integer(1)`)\cr
@@ -119,6 +120,7 @@ nn_tokenizer_num = nn_module(
 #' @name nn_tokenizer_categ
 #' @description
 #' Tokenizes categorical features into a dense embedding.
+#' For an input of shape `(batch, n_features)` the output shape is `(batch, n_features, d_token)`.
 #' @param cardinalities (`integer()`)\cr
 #'   The number of categories for each feature.
 #' @param d_token (`integer(1)`)\cr