## ---- include = FALSE---------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

## ----setup--------------------------------------------------------------------
library(luz)
library(torch)

## ---- eval = FALSE------------------------------------------------------------
#  net <- nn_module(
#    "Net",
#    initialize = function(num_class) {
#      self$conv1 <- nn_conv2d(1, 32, 3, 1)
#      self$conv2 <- nn_conv2d(32, 64, 3, 1)
#      self$dropout1 <- nn_dropout2d(0.25)
#      self$dropout2 <- nn_dropout2d(0.5)
#      self$fc1 <- nn_linear(9216, 128)
#      self$fc2 <- nn_linear(128, num_class)
#    },
#    forward = function(x) {
#      x <- self$conv1(x)
#      x <- nnf_relu(x)
#      x <- self$conv2(x)
#      x <- nnf_relu(x)
#      x <- nnf_max_pool2d(x, 2)
#      x <- self$dropout1(x)
#      x <- torch_flatten(x, start_dim = 2)
#      x <- self$fc1(x)
#      x <- nnf_relu(x)
#      x <- self$dropout2(x)
#      x <- self$fc2(x)
#      x
#    }
#  )

## ---- eval = FALSE------------------------------------------------------------
#  fitted <- net %>%
#    setup(
#      loss = nn_cross_entropy_loss(),
#      optimizer = optim_adam,
#      metrics = list(
#        luz_metric_accuracy
#      )
#    ) %>%
#    set_hparams(num_class = 10) %>%
#    set_opt_hparams(lr = 0.003) %>%
#    fit(train_dl, epochs = 10, valid_data = test_dl)

## ---- eval = FALSE------------------------------------------------------------
#  predictions <- predict(fitted, test_dl)

## ---- eval = FALSE------------------------------------------------------------
#  # -> Initialize objects: model, optimizers.
#  # -> Select fitting device.
#  # -> Move data, model, optimizers to the selected device.
#  # -> Start training
#  for (epoch in 1:epochs) {
#    # -> Training procedure
#    for (batch in train_dl) {
#      # -> Calculate model `forward` method.
#      # -> Calculate the loss
#      # -> Update weights
#      # -> Update metrics and tracking loss
#    }
#    # -> Validation procedure
#    for (batch in valid_dl) {
#      # -> Calculate model `forward` method.
#      # -> Calculate the loss
#      # -> Update metrics and tracking loss
#    }
#  }
#  # -> End training

## ---- eval=FALSE--------------------------------------------------------------
#  fitted <- net %>%
#    setup(
#      ...
#      metrics = list(
#        luz_metric_accuracy
#      )
#    ) %>%
#    fit(...)

## ---- eval = FALSE------------------------------------------------------------
#  luz_metric_accuracy <- luz_metric(
#    # An abbreviation to be shown in progress bars, or
#    # when printing progress
#    abbrev = "Acc",
#    # Initial setup for the metric. Metrics are initialized
#    # every epoch, for both training and validation
#    initialize = function() {
#      self$correct <- 0
#      self$total <- 0
#    },
#    # Run at every training or validation step and updates
#    # the internal state. The update function takes `preds`
#    # and `target` as parameters.
#    update = function(preds, target) {
#      pred <- torch::torch_argmax(preds, dim = 2)
#      self$correct <- self$correct + (pred == target)$
#        to(dtype = torch::torch_float())$
#        sum()$
#        item()
#      self$total <- self$total + pred$numel()
#    },
#    # Use the internal state to query the metric value
#    compute = function() {
#      self$correct/self$total
#    }
#  )

## ----include=FALSE, eval = torch::torch_is_installed()------------------------
library(luz)
torch::torch_manual_seed(1)
get_model <- function() {
  torch::nn_module(
    initialize = function(input_size, output_size) {
      self$fc <- torch::nn_linear(prod(input_size), prod(output_size))
      self$output_size <- output_size
    },
    forward = function(x) {
      out <- x %>%
        torch::torch_flatten(start_dim = 2) %>%
        self$fc()
      out$view(c(x$shape[1], self$output_size))
    }
  )
}

model <- get_model()
model <- model %>%
  setup(
    loss = torch::nn_mse_loss(),
    optimizer = torch::optim_adam,
    metrics = list(
      luz_metric_mae(),
      luz_metric_mse(),
      luz_metric_rmse()
    )
  ) %>%
  set_hparams(input_size = 10, output_size = 1) %>%
  set_opt_hparams(lr = 0.001)

x <- list(torch::torch_randn(100,10), torch::torch_randn(100, 1))

fitted <- model %>% fit(
  x,
  epochs = 1,
  verbose = FALSE,
  dataloader_options = list(batch_size = 2, shuffle = FALSE)
)

evaluation <- fitted %>% evaluate(data = x)

## ---- eval = FALSE------------------------------------------------------------
#  evaluation <- fitted %>% evaluate(data = valid_dl)
#  metrics <- get_metrics(evaluation)
#  print(evaluation)

## ----echo=FALSE, eval=torch::torch_is_installed()-----------------------------
options(cli.unicode = FALSE)
metrics <- get_metrics(evaluation)
print(evaluation)

## ---- eval = FALSE------------------------------------------------------------
#  print_callback <- luz_callback(
#    name = "print_callback",
#    initialize = function(message) {
#      self$message <- message
#    },
#    on_train_batch_end = function() {
#      cat("Iteration ", ctx$iter, "\n")
#    },
#    on_epoch_end = function() {
#      cat(self$message, "\n")
#    }
#  )

## ---- eval = FALSE------------------------------------------------------------
#  fitted <- net %>%
#    setup(...) %>%
#    fit(..., callbacks = list(
#      print_callback(message = "Done!")
#    ))