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:
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. However, by default, one pass will be conducted in âtrainâ mode, and another one in âevalâ mode, which is different from the PyTorch behavior. One can opt out of this by specifying respect_mode = FALSE
which will only trace the forward pass in the mode the network is currently in. 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))
Limitations of tracing
# 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))
ScriptModule
, operations that have different behaviors in training and eval modes will always behave as if it were in the mode it was in during tracing, no matter which mode the ScriptModule
is in. For example: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)))
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4