trainers.extensions.wandb_log_difftre

trainers.extensions.wandb_log_difftre#

wandb_log_difftre(run, trainer, plot_fns=None)[source]#

Logs DiffTRe training statistics to Weights & Biases.

Parameters:

  • run – Active W&B run

  • trainer (Difftre) – Trainer to log to W&B

  • plot_fns – Dictionary with functions to plot the selected predictions for selected statepoints

Example usage:

After initiating the trainer and the run, add a W&B tracking task to the trainer via:

def plot_fn(some_prediction):
    fig = plt.figure()
    ...
    return fig

plot_fns = {
    0: {'some_prediction': plot_fn}
}

wandb_log_difftre(run, difftre_trainer, plot_fn)
Contents