Base Class for Torch Module Constructor Wrappers

PipeOpTorch is the base class for all PipeOps that represent neural network layers in a Graph. During training, it generates a PipeOpModule that wraps an nn_module and attaches it to the architecture, which is also represented as a Graph consisting mostly of PipeOpModules an PipeOpNOPs.

The convenient way to construct such a PipeOp is the nn() helper, which prefixes the given key with "nn_" to look it up in the mlr_pipeops dictionary and uses the unprefixed key as the id of the resulting PipeOp: nn("linear", out_features = 10) is equivalent to po("nn_linear", id = "linear", out_features = 10). Because ids must be unique within a Graph, repeated layers can be disambiguated with a _<n> suffix, e.g. nn("linear_1") and nn("linear_2").

While the former Graph operates on ModelDescriptors, the latter operates on tensors.

The relationship between a PipeOpTorch and a PipeOpModule is similar to the relationshop between a nn_module_generator (like nn_linear) and a nn_module (like the output of nn_linear(...)). A crucial difference is that the PipeOpTorch infers auxiliary parameters (like in_features for nn_linear) automatically from the intermediate tensor shapes that are being communicated through the ModelDescriptor.

During prediction, PipeOpTorch takes in a Task in each channel and outputs the same new Task resulting from their feature union in each channel. If there is only one input and output channel, the task is simply piped through.

Parameters

The ParamSet is specified by the child class inheriting from PipeOpTorch. Usually the parameters are the arguments of the wrapped nn_module minus the auxiliary parameter that can be automatically inferred from the shapes of the input tensors.

Inheriting

When inheriting from this class, one should overload either the private$.shapes_out() and the private$.shape_dependent_params() methods, or overload private$.make_module().

.make_module(shapes_in, param_vals, task)
(list(), list(), Task or NULL) -> nn_module
This private method is called to generate the nn_module that is passed as argument module to PipeOpModule. It must be overwritten, when no module_generator is provided. If left as is, it calls the provided module_generator with the arguments obtained by the private method .shape_dependent_params().
.shapes_out(shapes_in, param_vals, task)
(list(), list(), Task or NULL) -> named list()
This private method gets a list of integer vectors (shapes_in), the parameter values (param_vals), as well as an (optional) Task. The shapes_in list indicates the shape of input tensors that will be fed to the module's $forward() function. The list has one item per input tensor, typically only one. The function should return a list of shapes of tensors that are created by the module. The shapes_in are named after the input channels of the PipeOp and are in the same order. The output shapes must be in the same order as the output names of the PipeOp. In case the output shapes depends on the task (as is the case for PipeOpTorchHead), the function should return valid output shapes (possibly containing NAs) whether or not the task argument is provided. Any dimension of shapes_in can be NA, i.e. unknown, so this method must not assume that a dimension it reads is known. It has to assert the dimensions it actually needs and propagate the NAs it can live with.

The inference should generally be permissive. For example, applying a convolutional layer to an input of shape c(NA, 3, NA, NA) should assume that the last two dimensions are big enough for the given kernel size. This can sometimes lead to runtime errors but this is preferable over rejecting valid architectures.

There are various assertion helpers for common input checks, such as assert_known_dims() and the "See also" links on its page. There are also shape_helpers, which provide the shape arithmetic (broadcasting, resolving negative dimension indices).
.shape_dependent_params(shapes_in, param_vals, task)
(list(), list(), Task or NULL) -> named list()
This private method has the same inputs as .shapes_out. If .make_module() is not overwritten, it constructs the arguments passed to module_generator. Usually this means that it must infer the auxiliary parameters that can be inferred from the input shapes and add it to the user-supplied parameter values (param_vals).

Input and Output Channels

During training, all inputs and outputs are of class ModelDescriptor. During prediction, all input and output channels are of class Task.

Shape Inference

A network is assembled without any data flowing through it, so mlr3torch tracks the shape of the tensors instead: it starts from the shape the ingress announces and hands each PipeOp the shapes of its inputs, which is how auxiliary parameters such as in_features of nn_linear are filled in automatically.

A shape is an integer() whose first entry is the batch dimension, e.g. c(NA, 3, 32, 32) for a batch of RGB images of 32 by 32 pixels. A dimension whose size is not known in advance is NA, and any dimension can be NA, not only the batch dimension: the sequence length of a transformer or the height and width of an image may equally well vary.

The public $shapes_out() method answers the same question for a single PipeOp, without building a network around it, which is the quickest way to see what an operator makes of a given input shape. It is also what reports the problem when an operator needs a dimension that is unknown, naming the PipeOp and the shape it was given.

Implementing a PipeOp that computes its own output shapes is described under .shapes_out() in the "Inheriting" section above.

State

The state is the value calculated by the public method shapes_out().

Internals

During training, the PipeOpTorch creates a PipeOpModule for the given parameter specification and the input shapes from the incoming ModelDescriptors using the private method .make_module(). The input shapes are provided by the slot pointer_shape of the incoming ModelDescriptors. The channel names of this PipeOpModule are identical to the channel names of the generating PipeOpTorch.

A model descriptor union of all incoming ModelDescriptors is then created. Note that this modifies the graph of the first ModelDescriptor in place for efficiency. The PipeOpModule is added to the graph slot of this union and the edges that connect the sending PipeOpModules to the input channel of this PipeOpModule are addeded to the graph. This is possible because every incoming ModelDescriptor contains the information about the id and the channel name of the sending PipeOp in the slot pointer.

The new graph in the model_descriptor_union represents the current state of the neural network architecture. It is structurally similar to the subgraph that consists of all pipeops of class PipeOpTorch and PipeOpTorchIngress that are ancestors of this PipeOpTorch.

For the output, a shallow copy of the ModelDescriptor is created and the pointer and pointer_shape are updated accordingly. The shallow copy means that all ModelDescriptors point to the same Graph which allows the graph to be modified by-reference in different parts of the code.

Super class

mlr3pipelines::PipeOp -> PipeOpTorch

Public fields

module_generator: (nn_module_generator or NULL)
The module generator wrapped by this PipeOpTorch. If NULL, the private method private$.make_module(shapes_in, param_vals) must be overwritte, see section 'Inheriting'. Do not change this after construction.

Methods

Inherited methods

`PipeOpTorch$new()`

Creates a new instance of this R6 class.

Usage

PipeOpTorch$new(
  id,
  module_generator,
  param_set = ps(),
  param_vals = list(),
  inname = "input",
  outname = "output",
  packages = "torch",
  tags = NULL
)

Arguments

id: (character(1))
Identifier of the resulting object.
module_generator: (nn_module_generator)
The torch module generator.
param_set: (ParamSet)
The parameter set.
param_vals: (list())
List of hyperparameter settings, overwriting the hyperparameter settings that would otherwise be set during construction.
inname: (character())
The names of the PipeOp's input channels. These will be the input channels of the generated PipeOpModule. Unless the wrapped module_generator's forward method (if present) has the argument ..., inname must be identical to those argument names in order to avoid any ambiguity.
If the forward method has the argument ..., the order of the input channels determines how the tensors will be passed to the wrapped nn_module.
If left as NULL (default), the argument module_generator must be given and the argument names of the modue_generator's forward function are set as inname.
outname: (character())
The names of the output channels. These will be the ouput channels of the generated PipeOpModule and therefore also the names of the list returned by its $train(). In case there is more than one output channel, the nn_module that is constructed by this PipeOp during training must return a named list(), where the names of the list are the names out the output channels. The default is "output".
packages: (character())
The R packages this object depends on.
tags: (character())
The tags of the PipeOp. The tags "torch" is always added.

`PipeOpTorch$shapes_out()`

Calculates the output shapes for the given input shapes, parameters and task.

Usage

PipeOpTorch$shapes_out(shapes_in, task = NULL)

Arguments

shapes_in: (list() of integer())
The input shapes, which must be in the same order as the input channel names of the PipeOp.
task: (Task or NULL)
The task, which is very rarely used (default is NULL). An exception is PipeOpTorchHead.

Returns

A named list() containing the output shapes. The names are the names of the output channels of the PipeOp.

`PipeOpTorch$clone()`

The objects of this class are cloneable with this method.

Usage

PipeOpTorch$clone(deep = FALSE)

Arguments

deep: Whether to make a deep clone.

Examples

## Creating a neural network
# In torch

task = tsk("iris")

network_generator = torch::nn_module(
  initialize = function(task, d_hidden) {
    d_in = length(task$feature_names)
    self$linear = torch::nn_linear(d_in, d_hidden)
    self$output = if (task$task_type == "regr") {
      torch::nn_linear(d_hidden, 1)
    } else if (task$task_type == "classif") {
      torch::nn_linear(d_hidden, output_dim_for(task))
    }
  },
  forward = function(x) {
    x = self$linear(x)
    x = torch::nnf_relu(x)
    self$output(x)
  }
)

network = network_generator(task, d_hidden = 50)
x = torch::torch_tensor(as.matrix(task$data(1, task$feature_names)))
y = torch::with_no_grad(network(x))


# In mlr3torch
network_generator = po("torch_ingress_num") %>>%
  nn("linear", out_features = 50) %>>%
  nn("head")
md = network_generator$train(task)[[1L]]
network = model_descriptor_to_module(md)
y = torch::with_no_grad(network(torch_ingress_num.input = x))



## Implementing a custom PipeOpTorch

# defining a custom module
nn_custom = nn_module("nn_custom",
  initialize = function(d_in1, d_in2, d_out1, d_out2, bias = TRUE) {
    self$linear1 = nn_linear(d_in1, d_out1, bias)
    self$linear2 = nn_linear(d_in2, d_out2, bias)
  },
  forward = function(input1, input2) {
    output1 = self$linear1(input1)
    output2 = self$linear1(input2)

    list(output1 = output1, output2 = output2)
  }
)

# wrapping the module into a custom PipeOpTorch

library(paradox)

PipeOpTorchCustom = R6::R6Class("PipeOpTorchCustom",
  inherit = PipeOpTorch,
  public = list(
    initialize = function(id = "nn_custom", param_vals = list()) {
      param_set = ps(
        d_out1 = p_int(lower = 1, tags = c("required", "train")),
        d_out2 = p_int(lower = 1, tags = c("required", "train")),
        bias = p_lgl(default = TRUE, tags = "train")
      )
      super$initialize(
        id = id,
        param_vals = param_vals,
        param_set = param_set,
        inname = c("input1", "input2"),
        outname = c("output1", "output2"),
        module_generator = nn_custom
      )
    }
  ),
  private = list(
    .shape_dependent_params = function(shapes_in, param_vals, task) {
      c(param_vals,
        list(d_in1 = tail(shapes_in[["input1"]], 1)), d_in2 = tail(shapes_in[["input2"]], 1)
      )
    },
    .shapes_out = function(shapes_in, param_vals, task) {
      list(
        input1 = c(head(shapes_in[["input1"]], -1), param_vals$d_out1),
        input2 = c(head(shapes_in[["input2"]], -1), param_vals$d_out2)
      )
    }
  )
)

## Training

# generate input
task = tsk("iris")
task1 = task$clone()$select(paste0("Sepal.", c("Length", "Width")))
task2 = task$clone()$select(paste0("Petal.", c("Length", "Width")))
graph = gunion(list(po("torch_ingress_num_1"), po("torch_ingress_num_2")))
mds_in = graph$train(list(task1, task2), single_input = FALSE)

mds_in[[1L]][c("graph", "task", "ingress", "pointer", "pointer_shape")]
#> $graph
#> 
#> ── Graph with 1 PipeOps: ───────────────────────────────────────────────────────
#>                   ID         State sccssors prdcssors
#>               <char>        <char>   <char>    <char>
#>  torch_ingress_num_1 <<UNTRAINED>>                   
#> 
#> ── Pipeline: <INPUT> -> torch_ingress_num_1 -> <OUTPUT> 
#> 
#> $task
#> 
#> ── <TaskClassif> (150x3): Iris Flowers ─────────────────────────────────────────
#> • Target: Species
#> • Properties: multiclass
#> • Features (2):
#>   • dbl (2): Sepal.Length, Sepal.Width
#> • Target classes: setosa (33%), versicolor (33%), virginica (33%)
#> 
#> $ingress
#> $ingress$torch_ingress_num_1.input
#> Ingress: Task[selector_name(c("Sepal.Length", "Sepal.Width"), assert_present = TRUE)] --> Tensor(NA, 2)
#> 
#> 
#> $pointer
#> [1] "torch_ingress_num_1" "output"             
#> 
#> $pointer_shape
#> [1] NA  2
#> 
mds_in[[2L]][c("graph", "task", "ingress", "pointer", "pointer_shape")]
#> $graph
#> 
#> ── Graph with 1 PipeOps: ───────────────────────────────────────────────────────
#>                   ID         State sccssors prdcssors
#>               <char>        <char>   <char>    <char>
#>  torch_ingress_num_2 <<UNTRAINED>>                   
#> 
#> ── Pipeline: <INPUT> -> torch_ingress_num_2 -> <OUTPUT> 
#> 
#> $task
#> 
#> ── <TaskClassif> (150x3): Iris Flowers ─────────────────────────────────────────
#> • Target: Species
#> • Properties: multiclass
#> • Features (2):
#>   • dbl (2): Petal.Length, Petal.Width
#> • Target classes: setosa (33%), versicolor (33%), virginica (33%)
#> 
#> $ingress
#> $ingress$torch_ingress_num_2.input
#> Ingress: Task[selector_name(c("Petal.Length", "Petal.Width"), assert_present = TRUE)] --> Tensor(NA, 2)
#> 
#> 
#> $pointer
#> [1] "torch_ingress_num_2" "output"             
#> 
#> $pointer_shape
#> [1] NA  2
#> 

# creating the PipeOpTorch and training it
po_torch = PipeOpTorchCustom$new()
po_torch$param_set$values = list(d_out1 = 10, d_out2 = 20)
train_input = list(input1 = mds_in[[1L]], input2 = mds_in[[2L]])
mds_out = do.call(po_torch$train, args = list(input = train_input))
po_torch$state
#> $output1
#> [1] NA 10
#> 
#> $output2
#> [1] NA 20
#> 

# the new model descriptors

# the resulting graphs are identical
identical(mds_out[[1L]]$graph, mds_out[[2L]]$graph)
#> [1] TRUE
# note that as a side-effect, also one of the input graphs is modified in-place for efficiency
mds_in[[1L]]$graph$edges
#>                 src_id src_channel    dst_id dst_channel
#>                 <char>      <char>    <char>      <char>
#> 1: torch_ingress_num_1      output nn_custom      input1
#> 2: torch_ingress_num_2      output nn_custom      input2

# The new task has both Sepal and Petal features
identical(mds_out[[1L]]$task, mds_out[[2L]]$task)
#> [1] TRUE
mds_out[[2L]]$task
#> 
#> ── <TaskClassif> (150x5): Iris Flowers ─────────────────────────────────────────
#> • Target: Species
#> • Properties: multiclass
#> • Features (4):
#>   • dbl (4): Petal.Length, Petal.Width, Sepal.Length, Sepal.Width
#> • Target classes: setosa (33%), versicolor (33%), virginica (33%)

# The new ingress slot contains all ingressors
identical(mds_out[[1L]]$ingress, mds_out[[2L]]$ingress)
#> [1] TRUE
mds_out[[1L]]$ingress
#> $torch_ingress_num_1.input
#> Ingress: Task[selector_name(c("Sepal.Length", "Sepal.Width"), assert_present = TRUE)] --> Tensor(NA, 2)
#> 
#> $torch_ingress_num_2.input
#> Ingress: Task[selector_name(c("Petal.Length", "Petal.Width"), assert_present = TRUE)] --> Tensor(NA, 2)
#> 

# The pointer and pointer_shape slots are different
mds_out[[1L]]$pointer
#> [1] "nn_custom" "output1"  
mds_out[[2L]]$pointer
#> [1] "nn_custom" "output2"  

mds_out[[1L]]$pointer_shape
#> [1] NA 10
mds_out[[2L]]$pointer_shape
#> [1] NA 20

## Prediction
predict_input = list(input1 = task1, input2 = task2)
tasks_out = do.call(po_torch$predict, args = list(input = predict_input))
identical(tasks_out[[1L]], tasks_out[[2L]])
#> [1] TRUE

## Shape inference

# `$shapes_out()` reports what an operator makes of an input shape, without building a network
conv = nn("conv2d", out_channels = 4, kernel_size = 3)
conv$shapes_out(list(c(NA, 3, 32, 32)))
#> $output
#> [1] NA  4 30 30
#> 

# any dimension may be unknown, so images of varying size work just as well
conv$shapes_out(list(c(NA, 3, NA, NA)))
#> $output
#> [1] NA  4 NA NA
#> 

# a dimension the operator does need is reported, naming the PipeOp and the shape it was given
try(conv$shapes_out(list(c(NA, NA, 32, 32))))
#> Error : PipeOp 'conv2d' requires the channel dimension (dimension 2) of the input shape to be known, but got shape (NA,NA,32,32).

Parameters

Inheriting

Input and Output Channels

Shape Inference

State

Internals

See also

Super class

Public fields

Methods

Public methods

PipeOpTorch$new()

Usage

Arguments

PipeOpTorch$shapes_out()

Usage

Arguments

Returns

PipeOpTorch$clone()

Usage

Arguments

Examples

`PipeOpTorch$new()`

`PipeOpTorch$shapes_out()`

`PipeOpTorch$clone()`