TorchScript

Tracing

TorchScript programs can be created from R using tracing. When using tracing, code is automatically converted into this subset of Python by recording only the actual operators on tensors and simply executing and discarding the other surrounding R code.

Currently tracing is the only supported way to create TorchScript programs from R code.

For example, let’s use the jit_trace function to create a TorchScript program. We pass a regular R function and example inputs.

fn <- function(x) {
  torch_relu(x)
}

traced_fn <- jit_trace(fn, torch_tensor(c(-1, 0, 1)))

The jit_trace function has executed the R function with the example input and recorded all torch operations that occurred during execution to create a graph. graph is how we call the intermediate representation of TorchScript programs, and it can be inspected with:

traced_fn$graph

The traced function can now be invoked as a regular R function:

traced_fn(torch_randn(3))

It’s also possible to trace nn_modules() defined in R, for example:

module <- nn_module(
  initialize = function() {
    self$linear1 <- nn_linear(10, 10)
    self$linear2 <- nn_linear(10, 1)
  },
  forward = function(x) {
    x %>% 
      self$linear1() %>% 
      nnf_relu() %>% 
      self$linear2()
  }
)
traced_module <- jit_trace(module(), torch_randn(10, 10))

When using jit_trace with a nn_module only the forward method is traced. You can use the jit_trace_module function to pass example inputs to other methods. Traced modules look like normal nn_modules(), and can be called the same way:

traced_module(torch_randn(3, 10))

# fn does does an operation for each dimension of a tensor
fn <- function(x) {
  x %>% 
    torch_unbind(dim = 1) %>% 
    lapply(function(x) x$sum()) %>% 
    torch_stack(dim = 1)
}
# we trace using as an example a tensor with size (10, 5, 5)
traced_fn <- jit_trace(fn, torch_randn(10, 5, 5))
# applying it with a tensor with different size returns an error.
traced_fn(torch_randn(11, 5, 5))

traced_dropout <- jit_trace(nn_dropout(), torch_ones(5,5))
traced_dropout(torch_ones(3,3))
traced_dropout$eval()
# even after setting to eval mode, dropout is applied
traced_dropout(torch_ones(3,3))

fn <- function(x, y) {
  x + y
}
jit_trace(fn, torch_tensor(1), 1)

Compiling TorchScript

It’s also possible to create TorchScript programs by compiling TorchScript code. TorchScript code looks a lot like standard python code. For example:

tr <- jit_compile("
def fn (x: Tensor):
  return torch.relu(x)

")
tr$fn(torch_tensor(c(-1, 0, 1)))