luz metrics do not work with optim_lbfgs.

[wbzyl]

The code below produces errors, reproduced below, when the line with metrics is uncommented.

optimizer = optim_lbfgs,
metrics = list(luz_metric_mae())

The errors:

Epoch 1/20
Error in `FUN()`:
! Error while calling callback with class <metrics_callback/LuzCallback/R6> at
  on_train_batch_end.
Caused by error in `FUN()`:
! Error when evaluating update for metric with abbrev "MAE" and class <LuzMetric/R6>
ℹ The error happened at iter 1 of epoch 1.
ℹ The model was in training mode.
Caused by error in `torch_tensor_cpp()`:
! R type not handled

This is the code that produces errors.

#! /usr/bin/env Rscript

library(torch)
library(luz)

n <- 1000
d_in <- 3

x <- torch_randn(n, d_in)
coefs <- c(0.2, -1.3, -0.5)
y <- x$matmul(coefs)$unsqueeze(2) + torch_randn(n, 1)

ds <- tensor_dataset(x, y)
dl <- dataloader(ds, batch_size = 100, shuffle = TRUE)

d_hidden <- 32
d_out <- 1

net <- nn_module(
  initialize = function(d_in, d_hidden, d_out) {
    self$net <- nn_sequential(
      nn_linear(d_in, d_hidden),
      nn_relu(),
      nn_linear(d_hidden, d_out)
    )
  },
  forward = function(x) {
    self$net(x)
  },
  step = function() {
    if (ctx$training) {
      closure <- function() {
        pred <- ctx$model(ctx$input)
        loss <- ctx$loss_fn(pred, ctx$target)
        loss$backward()
        loss
      }
      ctx$loss <- ctx$opt$step(closure)
    } else {
      pred <- ctx$model(ctx$input)
      ctx$loss <- ctx$loss_fn(pred, ctx$target)
    }
  }
)

train_ids <- sample(
  1:length(ds),
  size = 0.6 * length(ds))

valid_ids <- sample(
  setdiff(1:length(ds), train_ids),
  size = 0.2 * length(ds)
)

test_ids <- setdiff(1:length(ds), union(train_ids, valid_ids))

train_ds <- dataset_subset(ds, indices = train_ids)
valid_ds <- dataset_subset(ds, indices = valid_ids)
test_ds <-  dataset_subset(ds, indices = test_ids)

train_dl <- dataloader(train_ds, batch_size = 100, shuffle = TRUE)
valid_dl <- dataloader(valid_ds, batch_size = 100)
test_dl <-  dataloader(test_ds, batch_size = 100)

fitted <- net |>
  setup(
    loss = nn_mse_loss(),
    optimizer = optim_lbfgs,
    # metrics = list(luz_metric_mae())
  ) |>
  set_hparams(d_in = d_in, d_hidden = d_hidden, d_out = d_out) |>
  set_opt_hparams(line_search_fn = "strong_wolfe") |>
  fit(train_dl, epochs = 20, valid_data = valid_dl)

[dfalbel]

The following should work. In your custom step() function you must assign the predicted valeus to ctx$pred so they can be found by luz callback that computes the metrics.

net <- nn_module(
  initialize = function(d_in, d_hidden, d_out) {
    self$net <- nn_sequential(
      nn_linear(d_in, d_hidden),
      nn_relu(),
      nn_linear(d_hidden, d_out)
    )
  },
  forward = function(x) {
    self$net(x)
  },
  step = function() {
    if (ctx$training) {
      closure <- function() {
        ctx$pred <- ctx$model(ctx$input)
        loss <- ctx$loss_fn(ctx$pred, ctx$target)
        loss$backward()
        loss
      }
      ctx$loss <- ctx$opt$step(closure)
    } else {
      ctx$pred <- ctx$model(ctx$input)
      ctx$loss <- ctx$loss_fn(ctx$pred, ctx$target)
    }
  }
)

[wbzyl]

Now that the predicted values are found by a luz callback, the code works.

Thanks.