Adap test and way to get the current loss

mlverse · Dec 12, 2023 · 0da9c78 · 0da9c78
1 parent c8a23d7
commit 0da9c78
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/R/callbacks.R b/R/callbacks.R
@@ -436,6 +436,7 @@ luz_callback_lr_scheduler <- luz_callback(
     }
     self[[call_on]] <- function() {
       if ("metrics" %in% names(formals(self$scheduler$step))) {
+        current_loss <- ctx$loss[[self$opt_name]]
         self$scheduler$step(current_loss)
       } else {
         self$scheduler$step()

diff --git a/tests/testthat/_snaps/callbacks.md b/tests/testthat/_snaps/callbacks.md
@@ -13,6 +13,20 @@
       Adjusting learning rate of group 1 to 0.0001
       Adjusting learning rate of group 1 to 0.0000
 
+---
+
+    Code
+      expect_message({
+        output <- mod %>% set_hparams(input_size = 10, output_size = 1) %>% fit(dl,
+          verbose = FALSE, epochs = 20, callbacks = list(luz_callback_lr_scheduler(
+            torch::lr_reduce_on_plateau, verbose = TRUE, patience = 2, threshold = 0.1)))
+      })
+    Message
+      Epoch 7: reducing learning rate of group 1 to 1.0000e-05
+      Epoch 10: reducing learning rate of group 1 to 1.0000e-06
+      Epoch 13: reducing learning rate of group 1 to 1.0000e-07
+      Epoch 16: reducing learning rate of group 1 to 1.0000e-08
+
 # progressbar appears with training and validation
 
     Code

diff --git a/tests/testthat/test-callbacks.R b/tests/testthat/test-callbacks.R
@@ -29,8 +29,13 @@ test_that("callback lr scheduler", {
     expect_message({
       output <- mod %>%
         set_hparams(input_size = 10, output_size = 1) %>%
-        fit(dl, verbose = FALSE, epochs = 5, callbacks = list(
-          luz_callback_lr_scheduler(torch::lr_reduce_on_plateau, verbose = TRUE)
+        fit(dl, verbose = FALSE, epochs = 20, callbacks = list(
+          luz_callback_lr_scheduler(
+            torch::lr_reduce_on_plateau,
+            verbose = TRUE,
+            patience = 2,
+            threshold = 1e-1
+          )
         ))
     })
   })