ptyrax.logger

ptyrax.logger#

Functions

log_image(writer, tag, tensor, step, **kwargs)

Log an image to tensorboard.

log_on_batch_end(writer, current_losses, ...)

Log scalar loss and optional model-level batch metrics to TensorBoard.

log_on_epoch_end(writer, epoch, model, ...)

Hook executed at the end of each epoch to log predictions, scalars and other metrics.

log_on_train_start(writer, dataset[, debug])

Hook to log initial information at the start of training to TensorBoard.

tensorboard_to_hdf5(tb_dir, out_filepath)

Convert TensorBoard event files to a single HDF5 archive.
ptyrax.logger.log_image(writer, tag, tensor, step, **kwargs)[source]#

Log an image to tensorboard.

Parameters:

  • writer (SummaryWriter) – SummaryWriter Tensorboard writer

  • tag (str) – str Tag for the image

  • tensor (Inexact[Array, '...']) – np.ndarray Image to log

  • step (jaxtyping.Integer) – int Current step

  • **kwargs – Additional arguments for the plot function

Returns:

None

Return type:

None

ptyrax.logger.log_on_batch_end(writer, current_losses, model, step)[source]#

Log scalar loss and optional model-level batch metrics to TensorBoard.

Parameters:

  • writer (SummaryWriter) – TensorBoard SummaryWriter instance.

  • current_losses (float) – Loss value for the current batch.

  • model (ImagePredictionModel) – Model (checked for a __log_batch__ hook).

  • step (int) – Global step counter.

Return type:

None

ptyrax.logger.log_on_epoch_end(writer, epoch, model, epoch_total_loss, all_epoch_losses, debug=False, diff_pat_gamma=0.5, log_every=1, **kwargs)[source]#

Hook executed at the end of each epoch to log predictions, scalars and other metrics.

Parameters:

  • writer (SummaryWriter) – TensorBoard SummaryWriter.

  • epoch (int) – Current epoch index.

  • model (ImagePredictionModel) – Model used for prediction/logging.

  • epoch_total_loss (float) – Total loss for the epoch.

  • all_epoch_losses (dict) – Iterable of NamedLoss-like objects.

  • debug (bool) – If True, log per-sample predictions.

  • diff_pat_gamma (float) – Gamma applied to diffraction patterns when visualizing.

  • log_every (int) – Only run logging every log_every epochs.

Returns:

None

Return type:

None

ptyrax.logger.log_on_train_start(writer, dataset, debug=False)[source]#

Hook to log initial information at the start of training to TensorBoard.

Parameters:

  • writer (SummaryWriter) – TensorBoard SummaryWriter instance.

  • dataset (ImageDataset) – Dataset to inspect and optionally log example images from.

  • debug (bool) – If True, log additional debug images for every dataset item.

Returns:

None

Return type:

None

ptyrax.logger.tensorboard_to_hdf5(tb_dir, out_filepath)[source]#

Convert TensorBoard event files to a single HDF5 archive.

Reads scalars, histograms, text tensors, and images from the TensorBoard log directory and writes them into a structured HDF5 file for archival and post-hoc analysis.

Parameters:

  • tb_dir (str | Path) – Path to the directory containing TensorBoard event files.

  • out_filepath (str | Path) – Destination HDF5 file path.

Return type:

None

Contents