Repository: netket/netket

Issue:
Early stop when energy does not get any lower after some number of iterations
Hi, I'd like to know if this feature is already implemented somewhere. When sweeping hamiltonian parameters or netket's hyperparameters this feature would be handy to avoid long run times.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/netket/netket/pull/489?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`v3.0@5d0ccb5`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/netket/netket/pull/489/graphs/tree.svg?width=650&height=150&src=pr&token=gzcOlpO5lB)](https://codecov.io/gh/netket/netket/pull/489?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##             v3.0     #489   +/-   ##
=======================================
  Coverage        ?   51.71%           
=======================================
  Files           ?      103           
  Lines           ?     5482           
  Branches        ?      819           
=======================================
  Hits            ?     2835           
  Misses          ?     2443           
  Partials        ?      204           
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/netket/netket/pull/489?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/netket/netket/pull/489?src=pr&el=footer). Last update [5d0ccb5...d92e3a3](https://codecov.io/gh/netket/netket/pull/489?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

VolodyaCO: @PhilipVinc do you want me to write an abstract class with a callback method? (I shall change the earlystopping method name to just callback)
VolodyaCO: @PhilipVinc do I need to add a test to increase the coverage?
PhilipVinc: yes please
VolodyaCO: Adding a timeout can also be useful. @PhilipVinc do you want me to include it in this PR?
PhilipVinc: If you want, why not.
Il 20 ott 2020, 20:55 +0200, Vladimir Vargas-Calderón <notifications@github.com>, ha scritto:
> Adding a timeout can also be useful. @PhilipVinc do you want me to include it in this PR?
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub, or unsubscribe.

PhilipVinc: Sorry, probably I wasn't clear...
But I'd like to avoid polluting the run interface which already has lots of inputs.

The timeout can be easily implemented through a callback, so that's how I'd do it.
VolodyaCO: Ok, but it needs the tqdm object as an input. Do you want me to include the timeout as a separate callback from `EarlyStopping`? If so, the `run` method would receive a list of callbacks, though I have to see how to pass different arguments to each callback in the list...
PhilipVinc: I need to think about this.

I would not want to add yet another argument to `vmc.run`. there are already a lot, and this should be doable only with callbacks.

I mean, to zeroth-order you can play the same trick played with `best_cost_function`: initialise it to `NaN` and if it's nan simply get the time from `time.now()`, and then check compared to that one...
PhilipVinc: Yes the list of callbacks is what we want.
No, all callbacks should accept a fixed, pre-defined list of arguments.
we can discuss on what they should be, but it should not be too dynamical.
VolodyaCO: Yeah, I guess we can leave the same arguments you proposed, and implement
the timeout as a callback only using the step argument to control the
timeout.

On Fri, 23 Oct 2020 at 08:29, Filippo Vicentini <notifications@github.com>
wrote:

> Yes the list of callbacks is what we want.
> No, all callbacks should accept a fixed, pre-defined list of arguments.
> we can discuss on what they should be, but it should not be too dynamical.
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/netket/netket/pull/489#issuecomment-715341231>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AHQJA7353QONLNKQVXKA2ADSMGAL3ANCNFSM4SUYTR7Q>
> .
>


-- 
Best wishes,

Vladimir Vargas-Calderón
<https://www.researchgate.net/profile/Vladimir_Vargas-Calderon>
PhD Physics Student @ Universidad Nacional de Colombia
+57 318 407 9118

-- 
*Aviso legal:* El contenido de este mensaje y los archivos adjuntos son 
confidenciales y de uso exclusivo de la Universidad Nacional de Colombia. 
Se encuentran dirigidos sólo para el uso del destinatario al cual van 
enviados. La reproducción, lectura y/o copia se encuentran prohibidas a 
cualquier persona diferente a este y puede ser ilegal. Si usted lo ha 
recibido por error, infórmenos y elimínelo de su correo. Los Datos 
Personales serán tratados conforme a la Ley 1581 de 2012 y a nuestra 
Política de Datos Personales que podrá consultar en la página web 
www.unal.edu.co <http://www.unal.edu.co/>.* *Las opiniones, informaciones, 
conclusiones y cualquier otro tipo de dato contenido en este correo 
electrónico, no relacionados con la actividad de la Universidad Nacional de 
Colombia, se entenderá como personales y de ninguna manera son avaladas por 
la Universidad.

VolodyaCO: I have added the timeout as another callback, and now the vmc run method receives a callbacks list.
PhilipVinc: Great!
As a last thing, could you please make sure we accept both callbacks and lists of callbkacs? 

Internally just use the list, and do the conversion by checking if the provided callback does not have an __iter__method. See loggers for an example.
PhilipVinc: I think you only need to fix the tests to use the kwarg 'callback' instead of 'callbacks', fix the docstring to also use the singular (and specify that it can either be a single callback or an iterable of callbacks) and the example, and we're done!


Patch:
diff --git a/Examples/EarlyStopping/heisenberg1d.py b/Examples/EarlyStopping/heisenberg1d.py
new file mode 100644
index 00000000..5511e8ed
--- /dev/null
+++ b/Examples/EarlyStopping/heisenberg1d.py
@@ -0,0 +1,57 @@
+# Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#    http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import netket as nk
+
+# TODO Help, nk.callbacks.EarlyStopping raises AttributeError: module 'netket' has no attribute 'callbacks'
+
+# 1D Lattice
+g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)
+
+# Hilbert space of spins on the graph
+# with total Sz equal to 0
+hi = nk.hilbert.Spin(s=0.5, graph=g, total_sz=0)
+
+# Heisenberg hamiltonian
+ha = nk.operator.Heisenberg(hilbert=hi)
+
+# Symmetric RBM Spin Machine
+ma = nk.machine.RbmSpin(alpha=1, hilbert=hi, symmetry=True)
+ma.init_random_parameters(seed=1234, sigma=0.01)
+
+# Metropolis Exchange Sampling
+# Notice that this sampler exchanges two neighboring sites
+# thus preservers the total magnetization
+sa = nk.sampler.MetropolisExchange(machine=ma)
+
+# Optimizer
+op = nk.optimizer.Sgd(ma, learning_rate=0.05)
+
+# Stochastic Reconfifugration
+sr = nk.optimizer.SR(ma, diag_shift=0.1)
+
+# Early Stopping
+es = nk.callbacks.EarlyStopping(patience=10)
+
+# Variational Monte Carlo
+gs = nk.Vmc(
+    hamiltonian=ha,
+    sampler=sa,
+    optimizer=op,
+    n_samples=1000,
+    n_discard=2,
+    sr=sr,
+)
+
+gs.run(out="test", n_iter=300, callback=es)
diff --git a/netket/__init__.py b/netket/__init__.py
index 5e66b485..cdf39a67 100644
--- a/netket/__init__.py
+++ b/netket/__init__.py
@@ -16,6 +16,7 @@ __all__ = [
     "dynamics",
     "exact",
     "graph",
+    "callbacks",
     "hilbert",
     "machine",
     "operator",
@@ -31,6 +32,7 @@ __all__ = [
 from . import (
     dynamics,
     exact,
+    callbacks,
     graph,
     hilbert,
     logging,
diff --git a/netket/abstract_variational_driver.py b/netket/abstract_variational_driver.py
index 220e840e..7c9bdc08 100644
--- a/netket/abstract_variational_driver.py
+++ b/netket/abstract_variational_driver.py
@@ -15,6 +15,20 @@ from tqdm import tqdm
 import warnings
 
 
+def _to_iterable(maybe_iterable):
+    """
+    _to_iterable(maybe_iterable)
+
+    Ensure the result is iterable. If the input is not iterable, it is wrapped into a tuple.
+    """
+    if hasattr(maybe_iterable, "__iter__"):
+        surely_iterable = maybe_iterable
+    else:
+        surely_iterable = (maybe_iterable,)
+
+    return surely_iterable
+
+
 # Note: to implement a new Driver (see also _vmc.py for an example)
 # If you want to inherit the nice interface of AbstractMCDriver, you should
 # subclass it, defining the following methods:
@@ -160,6 +174,7 @@ class AbstractVariationalDriver(abc.ABC):
         save_params_every=50,  # for default logger
         write_every=50,  # for default logger
         step_size=1,  # for default logger
+        callback=lambda *x: True,
     ):
         """
         Executes the Monte Carlo Variational optimization, updating the weights of the network
@@ -178,6 +193,7 @@ class AbstractVariationalDriver(abc.ABC):
             logger is provided)
             :step_size: Every how many steps should observables be logged to disk (default=1)
             :show_progress: If true displays a progress bar (default=True)
+            :callback: Callable or list of callable callback functions to stop training given a condition
         """
 
         if not isinstance(n_iter, numbers.Number):
@@ -200,14 +216,15 @@ class AbstractVariationalDriver(abc.ABC):
             # if out is a path, create an overwriting Json Log for output
             if isinstance(out, str):
                 loggers = (_JsonLog(out, "w", save_params_every, write_every),)
-            elif hasattr(out, "__iter__"):
-                loggers = out
             else:
-                loggers = (out,)
+                loggers = _to_iterable(out)
         else:
             loggers = tuple()
             show_progress = False
 
+        callbacks = _to_iterable(callback)
+        callback_stop = False
+
         with tqdm(
             self.iter(n_iter, step_size), total=n_iter, disable=not show_progress
         ) as itr:
@@ -223,6 +240,13 @@ class AbstractVariationalDriver(abc.ABC):
                 for logger in loggers:
                     logger(self.step_count, log_data, self.machine)
 
+                for callback in callbacks:
+                    if not callback(step, log_data, self):
+                        callback_stop = True
+
+                if callback_stop:
+                    break
+
         # flush at the end of the evolution so that final values are saved to
         # file
         for logger in loggers:
diff --git a/netket/callbacks/__init__.py b/netket/callbacks/__init__.py
new file mode 100644
index 00000000..9596ab7d
--- /dev/null
+++ b/netket/callbacks/__init__.py
@@ -0,0 +1,2 @@
+from .earlystopping import EarlyStopping
+from .timeout import Timeout
diff --git a/netket/callbacks/earlystopping.py b/netket/callbacks/earlystopping.py
new file mode 100644
index 00000000..bfaf606b
--- /dev/null
+++ b/netket/callbacks/earlystopping.py
@@ -0,0 +1,49 @@
+from numpy import infty
+
+
+class EarlyStopping:
+    """A simple callback to stop NetKet if there are no improvements in the training"""
+
+    def __init__(self, min_delta=0, patience=0, baseline=None):
+        """
+        Constructs a new EarlyStopping object that monitors whether a driver is improving
+        over optimisation epochs based on `driver._loss_name`.
+
+        Args:
+            min_delta: Minimum change in the monitored quantity to qualify as an improvement.
+            patience: Number of epochs with no improvement after which training will be stopped.
+            baseline: Baseline value for the monitored quantity. Training will stop if the driver
+                hits the baseline.
+        """
+        self.__min_delta = min_delta
+        self.__patience = patience
+        self.__baseline = baseline
+        self._best_val = infty
+        self.__best_iter = 0
+
+    def __call__(self, step, log_data, driver):
+        """
+        A boolean function that determines whether or not to stop training.
+
+        Args:
+            step: An integer corresponding to the step (iteration or epoch) in training.
+            log_data: A dictionary containing log data for training.
+            driver: A NetKet variational driver.
+
+        Returns:
+            A boolean. If True, training continues, else, it does not.
+        """
+        loss = log_data[driver._loss_name].mean.real
+        if loss <= self._best_val:
+            self._best_val = loss
+            self.__best_iter = step
+        if self.__baseline is not None:
+            if loss <= self.__baseline:
+                return False
+        if (
+            step - self.__best_iter >= self.__patience
+            and loss > self._best_val - self.__min_delta
+        ):
+            return False
+        else:
+            return True
diff --git a/netket/callbacks/timeout.py b/netket/callbacks/timeout.py
new file mode 100644
index 00000000..23e726b3
--- /dev/null
+++ b/netket/callbacks/timeout.py
@@ -0,0 +1,45 @@
+from numpy import infty
+import time
+
+
+class Timeout:
+    """A simple callback to stop NetKet after some time has passed."""
+
+    def __init__(self, timeout):
+        """
+        Constructs a new Timeout object that monitors whether a driver has been training
+        for more than a given timeout in order to hard stop training.
+
+        Args:
+            timeout: Number of seconds to wait before hard stopping training.
+        """
+        assert timeout > 0
+        self.__timeout = timeout
+        self.__init_time = None
+
+    def reset(self):
+        """Resets the initial time of the training"""
+        self.__init_time = None
+
+    def __call__(self, step, log_data, driver):
+        """
+        A boolean function that determines whether or not to stop training.
+
+        Args:
+            step: An integer corresponding to the step (iteration or epoch) in training.
+            log_data: A dictionary containing log data for training.
+            driver: A NetKet variational driver.
+
+        Returns:
+            A boolean. If True, training continues, else, it does not.
+
+        Note:
+            This callback does not make use of `step`, `log_data` nor `driver`.
+        """
+        if self.__init_time is None:
+            self.__init_time = time.time()
+        else:
+            print(time.time() - self.__init_time)
+            if time.time() - self.__init_time >= self.__timeout:
+                return False
+        return True
diff --git a/setup.py b/setup.py
index 5a1aa1a0..bd5ae6aa 100644
--- a/setup.py
+++ b/setup.py
@@ -12,6 +12,7 @@ setup(
         "netket.graph",
         "netket.hilbert",
         "netket.logging",
+        "netket.callbacks",
         "netket.machine",
         "netket.machine.density_matrix",
         "netket.sampler",


Test patch:
diff --git a/Test/Callbacks/test_callbacks.py b/Test/Callbacks/test_callbacks.py
new file mode 100644
index 00000000..b8a06482
--- /dev/null
+++ b/Test/Callbacks/test_callbacks.py
@@ -0,0 +1,48 @@
+import netket as nk
+import time
+
+SEED = 3141592
+
+
+def _run_vmc(callbacks, n_iter=20):
+    nk.random.seed(SEED)
+    g = nk.graph.Hypercube(length=8, n_dim=1)
+    hi = nk.hilbert.Spin(s=0.5, graph=g)
+
+    ma = nk.machine.RbmSpin(hilbert=hi, alpha=1)
+    ma.init_random_parameters(sigma=0.01, seed=SEED)
+
+    ha = nk.operator.Ising(hi, h=1.0)
+    sa = nk.sampler.MetropolisLocal(machine=ma)
+
+    op = nk.optimizer.Sgd(ma, learning_rate=0.1)
+
+    vmc = nk.Vmc(hamiltonian=ha, sampler=sa, optimizer=op, n_samples=500)
+    st = time.time()
+    vmc.run(n_iter, callback=callbacks)
+    runtime = time.time() - st
+    return vmc.step_count, runtime
+
+
+def test_earlystopping_with_patience():
+    patience = 10
+    es = nk.callbacks.EarlyStopping(patience=patience)
+    es._best_val = -1e6
+    step_value = _run_vmc([es])
+    assert step_value, runtime == patience
+
+
+def test_earlystopping_with_baseline():
+    baseline = -10
+    es = nk.callbacks.EarlyStopping(baseline=baseline)
+    _step_value, runtime = _run_vmc([es])
+
+
+def test_timeout():
+    timeout = 5
+    tout = nk.callbacks.Timeout(timeout=timeout)
+    step_value, runtime = _run_vmc([tout], 300)
+
+    # There is a lag in the first iteration of about 3 seconds
+    # But the timeout works!
+    assert abs(timeout - runtime) < 3