R/utils.R

In mlr3torch: Deep Learning with 'mlr3'

#' Auto Device
#'
#' Resolves the `device` parameter of a learner.
#' `"auto"` becomes `"cuda"` when a CUDA device is available and `"cpu"` otherwise; any other value
#' is returned unchanged, except that an explicit `"cuda"` (or `"cuda:<index>"`) without an available
#' CUDA device with that index is an error rather than a failure deep inside libtorch later on.
#'
#' @param device (`character(1)`)\cr
#'   The device. If not `NULL`, is returned as is.
#' @export
auto_device = function(device = NULL) {
  if (!is.null(device) && device == "auto") {
    device = if (cuda_is_available()) "cuda" else "cpu"
    lg$debug("Auto-detected device '%s'.", device)
  }
  if (!is.null(device) && startsWith(device, "cuda")) {
    if (!cuda_is_available()) {
      stopf("Device is set to '%s', but no CUDA device is available. Set `device` to 'cpu', or to 'auto' to select automatically.", device) # nolint
    }
    index = sub("^cuda:?", "", device)
    if (nzchar(index) && as.integer(index) >= cuda_device_count()) {
      stopf("Device is set to '%s', but only %i CUDA device(s) are available (indices start at 0).", device, cuda_device_count()) # nolint
    }
  }
  return(device)
}

running_on_mac = function() {
  Sys.info()["sysname"] == "Darwin"
}

inferps = function(fn, ignore = character(0), tags = "train", required = FALSE) {
  if (inherits(fn, "R6ClassGenerator")) {
    fn = get_init(fn)
  }
  if (is.null(fn)) {
    return(ps())
  }
  assert_function(fn)
  assert_character(ignore, any.missing = FALSE)
  ignore = union(ignore, "...")
  frm = formals(fn)
  frm = frm[names(frm) %nin% ignore]

  # an argument that has no default cannot be left unset, so it is tagged "required"
  frm_domains = lapply(frm, function(formal) {
    p_uty(tags = if (required && identical(formal, alist(x = )$x)) c(tags, "required") else tags)
  })

  do.call(paradox::ps, frm_domains)
}


# `null_ok` must be `FALSE` for parameters the operator cannot do without
make_check_vector = function(d, null_ok = TRUE) {
  crate(function(x) {
    if ((null_ok && is.null(x)) || test_integerish(x, any.missing = FALSE) && (length(x) %in% c(1, d))) { # nolint
      return(TRUE)
    }
    tmp = if (d == 1) "." else sprintf(" or %s.", d)
    sprintf("Must be an integerish vector of length 1%s", tmp)
    }, d, null_ok, .parent = topenv())
}

assert_inherits_classname = function(class_generator, classname) {
  assert_class(class_generator, "R6ClassGenerator")
  while (!is.null(class_generator)) {
    if (class_generator$classname == classname) {
      return(TRUE)
    }
    class_generator = class_generator$get_inherit()
  }
  stopf("R6ClassGenerator does not generate object that inherits from %s.", classname)
}

get_init = function(x) {
  get_method(x, "initialize")
}

# jarl-ignore unused_function: called from man-roxygen/learner_example.R, which jarl does not scan
default_task_id = function(learner) {
  task_id = get0("task_id", envir = parent.frame(), inherits = FALSE)
  if (!is.null(task_id)) {
    return(task_id)
  }
  if (inherits(learner, "LearnerTorchImage")) {
    stopf("Currently not available!")
  }
  switch(learner$task_type,
    classif = "iris",
    regr = "mtcars",
    stopf("No default task type.")
  )

}

# the method of the generator itself or, failing that, of the class it inherits from
get_method = function(x, name) {
  if (is.null(x) || is.null(x$public_methods)) return(NULL)
  if (exists(name, x$public_methods, inherits = FALSE)) return(x$public_methods[[name]])
  Recall(x$get_inherit(), name)
}

sample_input_from_shapes = function(shapes, n = 1L) {
  assert_list(shapes, types = "numeric", min.len = 1)
  assert_int(n)
  imap(shapes, function(shape, nm) {
    shape[1] = n
    invoke(torch_randn, .args = as.list(shape))
  })
}

load_col_info = function(name) {
  readRDS(system.file("col_info", paste0(name, ".rds"), package = "mlr3torch"))
}


test_equal_col_info = function(x, y) {
  nms = c("id", "type", "levels")
  if (!(test_permutation(colnames(x), nms) && test_permutation(colnames(y), nms))) {
    return(FALSE)
  }

  x = x[order(get("id"))]
  y = y[order(get("id"))]

  isTRUE(all.equal(x$id, y$id)) && isTRUE(all.equal(x$type, y$type)) &&
    all(pmap_lgl(list(x = x$levels, y = y$levels), function(x, y) isTRUE(all.equal(x, y))))
}


uniqueify = function(new, existing) {
  make.unique(c(existing, new), sep = "_")[length(existing) + seq_along(new)]
}

#' @rdname shape_helpers
#' @export
shape_to_str = function(x) {
  if (is.numeric(x) || is.logical(x)) { # single shape
    return(sprintf("(%s)", paste0(x, collapse = ",")))
  }
  assert(test_list(x) || is.null(x))
  if (is.null(x)) {
    return("(<unknown>)")
  }

  shapedescs = map_chr(x, function(y) {
    if (is.null(y)) {
      return("<unknown>")
    }
    paste0("(", paste(y, collapse = ",", recycle0 = TRUE), ")")
  })
  if (test_named(x)) {
    repr = paste0("[", paste(paste0(names(x), ": ", shapedescs), collapse = "; ", recycle0 = TRUE), "]")
    return(repr)
  }
  paste0("[",  paste(shapedescs, collapse = ";", recycle0 = TRUE), "]")
}

test_equal = function(x, y) {
  isTRUE(all.equal(x, y))
}


dataset_image = dataset("image_dataset",
  initialize = function(uris) {
    self$uris = uris
  },
  .getitem = function(x) {
    list(x = torchvision::transform_to_tensor(torchvision::base_loader(self$uris[x])))
  },
  .length = function() {
    length(self$uris)
  }
)

list_to_batch = function(tensors) {
  torch_cat(map(tensors, function(x) x$unsqueeze(1)), dim = 1L)
}

auto_cache_lazy_tensors = function(lts) {
  if (length(lts) <= 1L) {
    return(FALSE)
  }
  anyDuplicated(unlist(map_if(lts, function(x) length(x) > 0, function(x) dd(x)$dataset_hash))) > 0L
}

#' Replace the Head of a Network
#'
#' Replaces the last layer of a pretrained network with a fresh [`torch::nn_linear`] that has
#' `d_out` output features, so a network trained on some other task can be fine-tuned on this one.
#' The new layer's input size is read off the layer it replaces, and its weights are newly
#' initialized while the rest of the network keeps its pretrained weights.
#' @param network ([`torch::nn_module`])\cr
#'   The network
#' @param d_out (`integer(1)`)\cr
#'   The number of output classes.
#' @export
#' @keywords internal
replace_head = function(network, d_out) {
  UseMethod("replace_head")
}

# Creates a check function (as expected by `custom_check`) that asserts that `x` inherits from
# `classes`. Named `make_check_class()` and not `check_class()` to avoid masking
# `checkmate::check_class()`.
make_check_class = function(classes) {
  assert_character(classes, any.missing = FALSE, min.len = 1L)
  crate(function(x) {
    check_class(x, classes)
  }, classes, .parent = topenv())
}

check_nn_module = function(x) {
  check_class(x, "nn_module")
}

check_nn_module_generator = function(x) {
  check_class(x, "nn_module_generator")
}

LossNone = function() {
  structure(list(), class = "LossNone")
}

OptimizerNone = function() {
  structure(list(), class = "OptimizerNone")
}

CallbacksNone = function() {
  structure(list(), class = "CallbacksNone")
}

get_example_batch = function(dl) {
  ds = dl$dataset
  if (length(ds) < 2) {
    stopf("Dataset needs to contain at least 2 observations")
  }
  if (!is.null(ds$.getbatch)) {
    ds$.getbatch(1:2)
  } else {
    torch_stack(list(
      ds$.getitem(1),
      ds$.getitem(2)
    ))
  }
}

# l: list containing args
# f: function to be called
# returns: l reordered so it can be passed by position
order_named_args = function(f, l) {
  args = formalArgs(f)
  x = match("...", args)
  f2 = f
  body(f2) = quote(as.list(match.call())[-1L])
  l2 = do.call(f2, l)
  # (function(..., x = 1) {})(1, 2, 3) works
  # (function(..., x = 1) {})(1, 2, x = 3) DOES NOT WORK
  if (!is.null(names(l2)) && !is.na(x) && x != length(args)) {
    stopf("Because arguments are passed to tracer by position, `...` must either be the only or the last argument when named arguments are also passed.")
  }
  l2
}



#' @title Network Output Dimension
#' @description
#' Calculates the output dimension of a neural network for a given task that is expected by
#' \CRANpkg{mlr3torch}.
#' For classification, this is the number of classes (unless it is a binary classification task,
#' where it is 1). For regression, it is 1.
#'
#' This is an S3 generic and the single place where \CRANpkg{mlr3torch} decides how many output neurons
#' a task needs: it is what [`PipeOpTorchHead`] and the [`LearnerTorch`]s that build their own head
#' ask. Adding a method for a new task type is therefore the way to support it, see the
#' "Supporting Other Task Types" section of [`PipeOpTorchHead`].
#'
#' @param x (any)\cr
#'   The task.
#' @param ... (any)\cr
#'   Additional arguments. Not used yet.
#' @return (`integer(1)`) The number of output neurons.
#' @seealso [`PipeOpTorchHead`]
#' @export
output_dim_for = function(x, ...) {
  UseMethod("output_dim_for")
}

#' @export
output_dim_for.TaskClassif = function(x, ...) {
  if ("twoclass" %in% x$properties) {
    return(1L)
  }
  length(x$class_names)
}

#' @export
output_dim_for.TaskRegr = function(x, ...) {
  1L
}

single_lazy_tensor = function(task) {
  identical(task$feature_types[, "type"][[1L]], "lazy_tensor")
}

n_num_features = function(task) {
  sum(task$feature_types$type %in% c("numeric", "integer"))
}

n_categ_features = function(task) {
  sum(task$feature_types$type %in% c("factor", "ordered", "logical"))
}

# Cardinalities of the categorical features of a task, in the order of `task$feature_names`, which
# is the column order of the batches that `po("torch_ingress_categ")` and `ingress_categ()` produce.
# Two things this must get right and that are easy to get wrong:
#  * `Task$levels()` returns `NULL` for `logical()` features, so their cardinality has to be
#    supplied explicitly (it is always 2). Taking `lengths(task$levels(...))` alone yields 0.
#  * `task$feature_names` and `task$feature_types` are not always in the same order (e.g. after
#    `po("scale")`, or when the feature column role is set explicitly), so the order must not be
#    taken from `task$feature_types`. Otherwise the cardinalities silently desync from the columns.
categ_cardinalities = function(task) {
  categ_types = c("factor", "ordered", "logical")
  features = intersect(task$feature_names, task$feature_types[get("type") %in% categ_types, get("id")])
  if (!length(features)) {
    return(integer(0))
  }
  types = task$feature_types[list(features), "type", on = "id"][[1L]]
  cardinalities = lengths(task$levels(features))[features]
  cardinalities[types == "logical"] = 2L
  set_names(as.integer(cardinalities), features)
}

rbind_arrays = function(xs) {
  d = dim(xs[[1L]])
  k = length(d)
  walk(xs, function(x) {
    if (!identical(dim(x)[-1L], d[-1L])) {
      stopf("Cannot combine arrays of dimensions (%s) and (%s), they differ beyond the first dimension.", paste(d, collapse = ", "), paste(dim(x), collapse = ", ")) # nolint
    }
  })
  if (k == 1L) {
    return(array(do.call(c, xs), dim = sum(map_int(xs, function(x) dim(x)[1L]))))
  }
  rotated = do.call(c, lapply(xs, function(x) aperm(x, c(2:k, 1L))))
  nrows = sum(map_int(xs, function(x) dim(x)[1L]))
  aperm(array(rotated, dim = c(d[-1L], nrows)), c(k, seq_len(k - 1L)))
}