Enhance Axon.Loop: Add patience option for checkpointing and implement related test case (#610)

Author: alvises

lib/axon/loop.ex

@@ -1147,7 +1147,6 @@ defmodule Axon.Loop do
         {status, state} = fun.(state)
         default = %{monitor => cur_criteria_value, :since_last_improvement => 0}
         updated_handler_meta = Map.put(handler_meta, name, default)
-
         {status, %{state | handler_metadata: updated_handler_meta}}
     end
   end
@@ -1221,6 +1220,7 @@ defmodule Axon.Loop do
         filter: :always,
         path: "checkpoint",
         file_pattern: &default_checkpoint_file/1,
+        patience: 3,
         mode: :min
       ])
 
@@ -1229,6 +1229,7 @@ defmodule Axon.Loop do
     {filter, opts} = Keyword.pop!(opts, :filter)
     {path, opts} = Keyword.pop!(opts, :path)
     {file_pattern, opts} = Keyword.pop!(opts, :file_pattern)
+    {patience, opts} = Keyword.pop!(opts, :patience)
     {mode, serialize_opts} = Keyword.pop!(opts, :mode)
 
     checkpoint_fun = &checkpoint_impl(&1, path, file_pattern, serialize_opts)
@@ -1237,7 +1238,8 @@ defmodule Axon.Loop do
       monitor(loop, criteria, checkpoint_fun, :checkpoint,
         mode: mode,
         event: event,
-        filter: filter
+        filter: filter,
+        patience: patience
       )
     else
       handle_event(loop, event, checkpoint_fun, filter)

test/axon/loop_test.exs

@@ -899,6 +899,51 @@ defmodule Axon.LoopTest do
       assert ["checkpoint_0_1.ckpt", "checkpoint_1_1.ckpt", "checkpoint_2_1.ckpt"] ==
                File.ls!("checkpoint_custom_path") |> Enum.sort()
     end
+
+    test "with :criteria, saves checkpoint after :patience is exceeded when metric stops improving",
+         %{loop: loop} do
+      loop
+      |> Loop.handle_event(:epoch_completed, fn
+        %{epoch: 0} = state ->
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(15))}
+          {:continue, state}
+
+        %{epoch: 1} = state ->
+          # loss is improved (15 -> 10)
+          # since_last_improvement = 0
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(10))}
+          {:continue, state}
+
+        %{epoch: 2} = state ->
+          # loss is improved (10 -> 5)
+          # since_last_improvement = 0
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(5))}
+          {:continue, state}
+
+        %{epoch: 3} = state ->
+          # loss is NOT improved (5 -> 5)
+          # since_last_improvement = 1
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(5))}
+          {:continue, state}
+
+        %{epoch: 4} = state ->
+          # loss is NOT improved (5 -> 5)
+          # since_last_improvement = 2
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(5))}
+          {:continue, state}
+
+        %{epoch: 5} = state ->
+          # loss is NOT improved (5 -> 5)
+          # since_last_improvement = 0 (goes back to 0 and waits for improvement)
+          state = %{state | metrics: Map.put(state.metrics, "loss", Nx.tensor(5))}
+          {:continue, state}
+      end)
+      |> Loop.checkpoint(criteria: "loss", mode: :min, patience: 2)
+      |> Loop.run([{Nx.tensor([[0]]), Nx.tensor([[1]])}], Axon.ModelState.empty(), epochs: 6)
+
+      # checkpoint_{epoch}_{iteration}.ckpt
+      assert ["checkpoint_5_1.ckpt"] == File.ls!("checkpoint")
+    end
   end
 
   describe "from_state" do