netket.logging.TensorBoardLog

class netket.logging.TensorBoardLog

Bases: AbstractCallback

Creates a tensorboard logger using tensorboardX’s summarywriter.

Refer to its documentation for further details

https://tensorboardx.readthedocs.io/en/latest/tensorboard.html

TensorBoardX must be installed.

This class is a full AbstractCallback and can be passed either as out=logger or inside the callbacks=[..., logger] list. When used as a callback the logger automatically captures the variational state snapshot taken just before the parameter update.

Parameters:

• logdir (string) – Save directory location. Default is runs/CURRENT_DATETIME_HOSTNAME, which changes after each run. Use hierarchical folder structure to compare between runs easily. e.g. pass in ‘runs/exp1’, ‘runs/exp2’, etc. for each new experiment to compare across them.

• comment (string) – Comment logdir suffix appended to the default logdir. If logdir is assigned, this argument has no effect.

• purge_step (int) – When logging crashes at step \(T+X\) and restarts at step \(T\), any events whose global_step larger or equal to \(T\) will be purged and hidden from TensorBoard. Note that crashed and resumed experiments should have the same logdir.

• max_queue (int) – Size of the queue for pending events and summaries before one of the ‘add’ calls forces a flush to disk. Default is ten items.

• flush_secs (int) – How often, in seconds, to flush the pending events and summaries to disk. Default is every two minutes.

• filename_suffix (string) – Suffix added to all event filenames in the logdir directory. More details on filename construction in tensorboard.summary.writer.event_file_writer.EventFileWriter.

• write_to_disk (boolean) – If pass False, TensorBoardLog will not write to disk.

• metadata (dict | None) – Optional flat dict of key/value pairs written as a JSON text summary at the start of the run (visible in the TensorBoard Text tab).

Tip

Use metadata to attach a flat dict of hyper-parameters (learning rate, system size, model type, …) to the run. They are written once as a JSON text entry and appear in the TensorBoard Text tab, keeping all run information in one place without external bookkeeping.

Examples

Logging optimisation to tensorboard.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> # create a summary writer with automatically generated folder name.
>>> writer = nk.logging.TensorBoardLog()
>>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/
>>> # create a summary writer using the specified folder name.
>>> writer = nk.logging.TensorBoardLog("my_experiment")
>>> # folder location: my_experiment
>>> # create a summary writer with comment appended.
>>> writer = nk.logging.TensorBoardLog(comment="LR_0.1_BATCH_16")
>>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/

Attaching metadata to record hyper-parameters.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> writer = nk.logging.TensorBoardLog(
...     "my_experiment",
...     metadata={"learning_rate": 0.01, "alpha": 1, "L": 20},
... )
>>> gs.run(n_iter=300, out=writer)
>>> # hyper-parameters appear in the TensorBoard Text tab as JSON

Using the logger as a callback.

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> import netket as nk
>>> writer = nk.logging.TensorBoardLog("my_experiment")
>>> gs.run(n_iter=300, callbacks=[writer])

Inheritance

Attributes

callback_order

Methods

__call__(step, item, variational_state=None)

Call self as a function.

Parameters:

• step (int)

• item (dict[str, Any])

• variational_state (VariationalState | None)

before_parameter_update(step, log_data, driver)

Called after all update logic has been computed and the step has been accepted, but before the driver applies the parameter update.

At this point:

• The loss and its gradient have been computed by compute_loss_and_update().

• The step has been accepted (not rejected by on_compute_update_end()).

• driver.step_count still refers to the current step — it has not yet been incremented.

• The variational state parameters have not yet changed.

This is the right place to estimate additional observables, add data to log_data, or take a snapshot of the state for logging. Callbacks with a lower callback_order run first, so observables callbacks (order 0) are guaranteed to populate log_data before logger callbacks (order 10) read it.

flush(variational_state=None)

Writes to file the content of this logger.

Parameters:

machine – optionally also writes the parameters of the machine.

on_compute_update_end(step, log_data, driver)

Callback called at the end of the compute update phase, after computing the loss and its gradient.

This is called before the parameters are updated, so it can be used to implement custom logic for rejecting a step based on the computed loss or gradient.

Return type:

bool

Returns:

A boolean indicating whether to reject the step (i.e. repeat it with the same parameters). If it returns None, it is treated as False.

on_compute_update_start(step, log_data, driver)

on_run_end(step, driver)

on_run_error(step, error, driver)

on_run_start(step, driver)

on_step_end(step, log_data, driver)

on_step_start(step, log_data, driver)

replace(**kwargs)

Replace the values of the fields of the object with the values of the keyword arguments. If the object is a dataclass, dataclasses.replace will be used. Otherwise, a new object will be created with the same type as the original object.

Return type:

TypeVar(P, bound= Pytree)

Parameters:

• self (P)

• kwargs (Any)