Source code for netket.driver.abstract_variational_driver

# Copyright 2021 The NetKet Authors - 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
import numbers
from functools import partial

from tqdm import tqdm

import jax
from jax.tree_util import tree_map

from netket.logging import JsonLog
from netket.utils import mpi

def _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
        surely_iterable = (maybe_iterable,)

    return surely_iterable

# Note: to implement a new Driver (see also for an example)
# If you want to inherit the nice interface of AbstractMCDriver, you should
# subclass it, defining the following methods:
# - Either _forward_and_backward or individually _forward, _backward, that should
#   compute the loss function and the gradient. If the driver is minimizing or
#   maximising some loss function, this quantity should be assigned to self._stats
#   in order to monitor it.
# - _estimate_stats should return the MC estimate of a single operator
# - reset should reset the driver (usually the sampler).
# - info should return a string with an overview of the driver.
# - The __init__ method should be called with the machine and the optimizer. If this
#   driver is minimising a loss function and you want it's name to show up automatically
#   in the progress bar/output files you should pass the optional keyword argument
#   minimized_quantity_name.
[docs]class AbstractVariationalDriver(abc.ABC): """Abstract base class for NetKet Variational Monte Carlo drivers""" def __init__(self, variational_state, optimizer, minimized_quantity_name=""): self._mynode = mpi.node_number self._mpi_nodes = mpi.n_nodes self._loss_stats = None self._loss_name = minimized_quantity_name self._step_count = 0 self._variational_state = variational_state self.optimizer = optimizer def _forward_and_backward(self): # pragma: no cover """ Performs the forward and backward pass at the same time. Concrete drivers should either override this method, or override individually _forward and _backward. Returns: the update for the weights. """ self._forward() dp = self._backward() return dp def _forward(self): """ Performs the forward pass, computing the loss function. Concrete should either implement _forward and _backward or the joint method _forward_and_backward. """ raise NotImplementedError() # pragma: no cover def _backward(self): """ Performs the backward pass, computing the update for the parameters. Concrete should either implement _forward and _backward or the joint method _forward_and_backward. """ raise NotImplementedError() # pragma: no cover def _estimate_stats(self, observable): """ Returns the MCMC statistics for the expectation value of an observable. Must be implemented by super-classes of AbstractVMC. :param observable: A quantum operator (netket observable) :return: """ return self.state.expect(observable)
[docs] def reset(self): """ Resets the driver. Concrete drivers should also call super().reset() to ensure that the step count is set to 0. """ self.state.reset() self._step_count = 0 pass
[docs] @abc.abstractmethod def info(self, depth=0): """ Returns an info string used to print information to screen about this driver. """ pass # pragma: no cover
@property def state(self): """ Returns the machine that is optimized by this driver. """ return self._variational_state @property def optimizer(self): """ The optimizer used to update the parameters at every iteration. """ return self._optimizer @optimizer.setter def optimizer(self, optimizer): self._optimizer = optimizer if optimizer is not None: self._optimizer_state = optimizer.init(self.state.parameters) @property def step_count(self): """ Returns a monotonic integer labelling all the steps performed by this driver. This can be used, for example, to identify the line in a log file. """ return self._step_count
[docs] def iter(self, n_steps: int, step: int = 1): """ Returns a generator which advances the VMC optimization, yielding after every `step_size` steps. Args: n_iter: The total number of steps to perform. step_size: The number of internal steps the simulation is advanced every turn. Yields: int: The current step. """ for _ in range(0, n_steps, step): for i in range(0, step): dp = self._forward_and_backward() if i == 0: yield self.step_count self._step_count += 1 self.update_parameters(dp)
[docs] def advance(self, steps: int = 1): """ Performs `steps` optimization steps. steps: (Default=1) number of steps """ for _ in self.iter(steps): pass
[docs] def run( self, n_iter, out=None, obs=None, show_progress=True, 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 stored in this driver for `n_iter` steps and dumping values of the observables `obs` in the output `logger`. If no logger is specified, creates a json file at `out`, overwriting files with the same prefix. By default uses :class:`nk.logging.JsonLog`. To know about the output format check it's documentation. The logger object is also returned at the end of this function so that you can inspect the results without reading the json output. Args: n_iter: the total number of iterations out: A logger object, or an iterable of loggers, to be used to store simulation log and data. If this argument is a string, it will be used as output prefix for the standard JSON logger. obs: An iterable containing all observables that should be computed save_params_every: Every how many steps the parameters of the network should be serialized to disk (ignored if logger is provided) write_every: Every how many steps the json data should be flushed to disk (ignored if 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): raise ValueError( "n_iter, the first positional argument to `run`, must be a number!" ) if obs is None: obs = {} if out is None: out = tuple() print( "No output specified (out=[apath|nk.logging.JsonLogger(...)])." "Running the optimization but not saving the output." ) # Log only non-root nodes if self._mynode == 0: # 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),) else: loggers = _to_iterable(out) else: loggers = tuple() show_progress = False callbacks = _to_iterable(callback) callback_stop = False with tqdm( total=n_iter, disable=not show_progress, dynamic_ncols=True, ) as pbar: old_step = self.step_count first_step = True for step in self.iter(n_iter, step_size): log_data = self.estimate(obs) self._log_additional_data(log_data, step) # if the cost-function is defined then report it in the progress bar if self._loss_stats is not None: pbar.set_postfix_str(self._loss_name + "=" + str(self._loss_stats)) log_data[self._loss_name] = self._loss_stats # Execute callbacks before loggers because they can append to log_data for callback in callbacks: if not callback(step, log_data, self): callback_stop = True for logger in loggers: logger(self.step_count, log_data, self.state) if len(callbacks) > 0: if mpi.mpi_any(callback_stop): break # Reset the timing of tqdm after the first step, to ignore compilation time if first_step: first_step = False pbar.unpause() # Update the progress bar pbar.update(self.step_count - old_step) old_step = self.step_count # Final update so that it shows up filled. pbar.update(self.step_count - old_step) # flush at the end of the evolution so that final values are saved to # file for logger in loggers: logger.flush(self.state) return loggers
[docs] def estimate(self, observables): """ Return MCMC statistics for the expectation value of observables in the current state of the driver. Args: observables: A pytree of operators for which statistics should be computed. Returns: A pytree of the same structure as the input, containing MCMC statistics for the corresponding operators as leaves. """ return tree_map(self._estimate_stats, observables)
[docs] def update_parameters(self, dp): """ Updates the parameters of the machine using the optimizer in this driver Args: dp: the pytree containing the updates to the parameters """ self._optimizer_state, self.state.parameters = apply_gradient( self._optimizer.update, self._optimizer_state, dp, self.state.parameters )
def _log_additional_data(self, log_dict, step): """ Method to be implemented in sub-classes of AbstractVariationalDriver to log additional data at every step. This method is called at every iteration when executing with `run`. Args: `log_dict`: The dictionary containing all logged data. It must be modified in-place adding new keys. `step`: the current step number. Returns: Nothing. The log dictionary should be modified in place. """ pass # pragma: no cover
@partial(jax.jit, static_argnums=0) def apply_gradient(optimizer_fun, optimizer_state, dp, params): import optax updates, new_optimizer_state = optimizer_fun(dp, optimizer_state, params) new_params = optax.apply_updates(params, updates) return new_optimizer_state, new_params