feat(nn_block): parameter trafo (#399)

Author: sebffischer

NEWS.md

@@ -16,6 +16,9 @@
 * feat: TabResNet learner now supports lazy tensors.
 * feat: The `LearnerTorch` base class now supports the private method `$.ingress_tokens(task, param_vals)`
   for generating the `torch::dataset`.
+* feat: `nn("block")` (which allows to repeat the same network segment multiple
+  times) now has an extra argument `trafo`, which allows to modify the
+  parameter values per layer.
 * feat: Shapes can now have multiple `NA`s and not only the batch dimension can be missing. However, most `nn()` operators still expect only one missing values and will throw an error if multiple dimensions are unknown.
 
 # mlr3torch 0.2.1

R/PipeOpTorchBlock.R

@@ -7,30 +7,46 @@
 #' `__<layer>`.
 #'
 #' @section Parameters:
-#' The parameters available for the block itself, as well as
+#' The parameters available for the provided `block`, as well as
 #' * `n_blocks` :: `integer(1)`\cr
 #'   How often to repeat the block.
+#' * `trafo` :: `function(i, param_vals, param_set) -> list()`\cr
+#'   A function that allows to transform the parameters vaues of each layer (`block`).
+#'   Here,
+#'   * `i` :: `integer(1)`\cr
+#'       is the index of the layer, ranging from `1` to `n_blocks`.
+#'   * `param_vals` :: named `list()`\cr
+#'       are the parameter values of the layer `i`.
+#'   * `param_set` :: [`ParamSet`][paradox::ParamSet]\cr
+#'       is the parameter set of the whole `PipeOpTorchBlock`.
+#'
+#'   The function must return the modified parameter values for the given layer.
+#'   This, e.g., allows for special behavior of the first or last layer.
 #' @section Input and Output Channels:
 #' The `PipeOp` sets its input and output channels to those from the `block` (Graph)
 #' it received during construction.
 #' @templateVar id nn_block
 #' @template pipeop_torch
 #' @export
 #' @examplesIf torch::torch_is_installed()
-#' block = po("nn_linear") %>>% po("nn_relu")
-#' po_block = po("nn_block", block,
-#' nn_linear.out_features = 10L, n_blocks = 3)
-#' network = po("torch_ingress_num") %>>%
-#' po_block %>>%
-#' po("nn_head") %>>%
-#' po("torch_loss", t_loss("cross_entropy")) %>>%
-#' po("torch_optimizer", t_opt("adam")) %>>%
-#' po("torch_model_classif",
-#'   batch_size = 50,
-#'   epochs = 3)
+#' # repeat a simple linear layer with ReLU activation 3 times, but set the bias for the last
+#' # layer to `FALSE`
+#' block = nn("linear") %>>% nn("relu")
 #'
-#' task = tsk("iris")
-#' network$train(task)
+#' blocks = nn("block", block,
+#'   linear.out_features = 10L, linear.bias = TRUE, n_blocks = 3,
+#'   trafo = function(i, param_vals, param_set) {
+#'     if (i  == param_set$get_values()$n_blocks) {
+#'       param_vals$linear.bias = FALSE
+#'     }
+#'     param_vals
+#'   })
+#' graph = po("torch_ingress_num") %>>%
+#'   blocks %>>%
+#'   nn("head")
+#' md = graph$train(tsk("iris"))[[1L]]
+#' network = model_descriptor_to_module(md)
+#' network
 PipeOpTorchBlock = R6Class("PipeOpTorchBlock",
   inherit = PipeOpTorch,
   public = list(
@@ -44,8 +60,12 @@ PipeOpTorchBlock = R6Class("PipeOpTorchBlock",
     initialize = function(block, id = "nn_block", param_vals = list()) {
       private$.block = as_graph(block)
       private$.param_set_base = ps(
-        n_blocks = p_int(lower = 0L, tags = c("train", "required"))
+        n_blocks = p_int(lower = 0L, tags = c("train", "required")),
+        trafo = p_uty(tags = "train", custom_check = crate(function(x) {
+          check_function(x, args = c("i", "param_vals", "param_set"))
+        }))
       )
+
       super$initialize(
         id = id,
         param_vals = param_vals,
@@ -68,11 +88,18 @@ PipeOpTorchBlock = R6Class("PipeOpTorchBlock",
   private = list(
     .block = NULL,
     .make_graph = function(block, n_blocks) {
+      trafo = self$param_set$get_values()$trafo
       graph = block
-      graph$update_ids(prefix = paste0(self$id, "."))
-      graphs = c(list(graph), replicate(n_blocks - 1L, graph$clone(deep = TRUE)))
+      graphs = c(replicate(n_blocks, graph$clone(deep = TRUE)))
+      if (!is.null(trafo)) {
+        param_vals = map(graphs, function(graph) graph$param_set$get_values())
+        walk(seq_along(param_vals), function(i) {
+          vals = trafo(i = i, param_vals = param_vals[[i]], param_set = self$param_set)
+          graphs[[i]]$param_set$values = vals
+        })
+      }
       lapply(seq_len(n_blocks), function(i) {
-        graphs[[i]]$update_ids(postfix = paste0("__", i))
+        graphs[[i]]$update_ids(prefix = paste0(self$id, "."), postfix = paste0("__", i))
       })
       Reduce(`%>>%`, graphs)
     },
@@ -112,10 +139,10 @@ PipeOpTorchBlock = R6Class("PipeOpTorchBlock",
       map(mdouts, "pointer_shape")
     },
     .train = function(inputs) {
-      if (self$param_set$values$n_blocks == 0L) {
+      param_vals = self$param_set$get_values()
+      if (param_vals$n_blocks == 0L) {
         return(inputs)
       }
-      param_vals = self$param_set$get_values(tags = "train")
       block = private$.block$clone(deep = TRUE)
       graph = private$.make_graph(block, param_vals$n_blocks)
       inputs = set_names(inputs, graph$input$name)
@@ -130,4 +157,4 @@ PipeOpTorchBlock = R6Class("PipeOpTorchBlock",
 
 
 #' @include aaa.R
-register_po("nn_block", PipeOpTorchBlock, metainf = list(block = as_graph(po("nop"))))
+register_po("nn_block", PipeOpTorchBlock, metainf = list(block = as_graph(po("nop"))))

R/nn.R

@@ -11,5 +11,9 @@
 #' # is the same as:
 #' po2 = nn("linear")
 nn = function(.key, ...) {
-  invoke(po, .obj = paste0("nn_", .key), .args = insert_named(list(id = .key), list(...)))
+  args = list(...)
+  if (is.null(args$id)) {
+    args$id = .key
+  }
+  invoke(po, .obj = paste0("nn_", .key), .args = args)
 }

man/mlr_pipeops_nn_block.Rd

@@ -94,4 +94,20 @@ test_that("0 blocks are possible", {
   md = po("torch_ingress_num")$train(list(tsk("iris")))[[1L]]
   mdout = nn("block", block = nn("linear", out_features = 10), n_blocks = 0)$train(list(md))[[1L]]
   expect_equal(mdout$pointer_shape, c(NA, 4))
-})
+})
+
+test_that("trafo works", {
+  graph = po("torch_ingress_num") %>>%
+    po("nn_block", nn("linear", out_features = 10),
+      n_blocks = 2L,
+      trafo = function(i, param_vals, param_set) {
+      if (i == 2) {
+        param_vals$linear.bias = FALSE
+      }
+      param_vals
+    })
+
+  network = model_descriptor_to_module(graph$train(tsk("iris"))[[1L]])
+  expect_class(network$module_list$`0`$bias, "torch_tensor")
+  expect_true(is.null(network$module_list$`1`$bias))
+})

tests/testthat/test_PipeOpTorchBlock.R

@@ -4,3 +4,16 @@ test_that("nn works", {
   expect_class(x, "PipeOpTorchLinear")
   expect_equal(x$param_set$values$out_features, 3)
 })
+
+test_that("overwrite id", {
+  obj = nn("linear", id = "abc")
+  expect_equal(obj$id, "abc")
+})
+
+test_that("unnamed arg", {
+  graph = po("torch_ingress_num") %>>% nn("block", nn("linear", out_features = 3), n_blocks = 2)
+  md = graph$train(tsk("iris"))[[1L]]
+  network = model_descriptor_to_module(md)
+  expect_equal(network$module_list[[1]]$out_features, 3)
+  expect_equal(network$module_list[[2]]$out_features, 3)
+})