# Plotting of the loss

**URL:** <https://discuss.bayesflow.org/t/plotting-of-the-loss/168>\
**Category:** General\
**Created:** [June 7, 2025, 11:48am UTC](https://discuss.bayesflow.org/t/plotting-of-the-loss/168 "2025-06-07T11:48:15Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![tdj2000](https://avatars.discourse-cdn.com/v4/letter/t/bbe5ce/32.png) [@tdj2000](https://discuss.bayesflow.org/u/tdj2000)\
**Post date:** [June 7, 2025, 11:48am UTC](https://discuss.bayesflow.org/t/plotting-of-the-loss/168/1 "2025-06-07T11:48:15Z")

</div>

Hi all,

I have been training some models using different summary networks. One of the things I am interested in is the decay of the loss function, especially in the first few epochs. However, when using f = bf.diagnostics.plots.loss(history, per\_training\_step=True), I no longer see all training steps during the first epoch and the value of the loss plotted for the first epoch does not match the value that is output in the training visual. Ideally, I would like the x-axis to start at 0, which I believe was also the case in BF version 1. Anyone else with the same issue? I have attached an image demonstrating what I mean.

 ![bfforum](https://canada1.discourse-cdn.com/flex007/uploads/bayesflow/original/1X/882cebdd8217bdc3e305f4ba5d8c9aa77c2b1da8.png)

---

<div class="post-metadata">

**Author:** ![valentin](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/valentin/32/103_2.png) [@valentin](https://discuss.bayesflow.org/u/valentin)\
**Post date:** [June 7, 2025, 12:29pm UTC](https://discuss.bayesflow.org/t/plotting-of-the-loss/168/2 "2025-06-07T12:29:00Z")

</div>

Hi,  
thanks for reaching out 🙂 As far I can tell the issue with the numbers not matching was caused by an issue that is already resolved on the `dev` branch, the fix will be in the next release.

I have taken a quick look at the code, and as far as I can tell the `per_training_step` functionality is currently broken/was removed, because the `history` object only contains per-epoch values. Sorry for that, and that the docstring is inaccurate there. Could you open an issue on that, so that we can check if we can/want to bring it back in version 2?

For now, I think there are two options to get the information you need. Either, if you use online training, you can use a small `num_batches_per_epoch` value, which will give you a better resolution, but might slow down training somewhat. Alternatively, you could use the following custom callback that saves the detailed history for you.

```python
class StepHistory(keras.callbacks.Callback):
    """Callback that records events into a `StepHistory` object in every batch."""

    def __init__ (self):
        super(). __init__ ()
        self.history = {}

    def on_train_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        for k, v in logs.items():
            self.history.setdefault(k, []).append(v)

    def on_epoch_end(self, epoch, logs=None):
        # Set the history attribute on the model after the epoch ends. This will
        # make sure that the state which is set is the latest one.
        self.model.step_history = self

```

You would pass this to `fit` like this: `history = workflow.fit_...(epochs=..., ..., callbacks=[StepHistory()])`. If you use the BasicWorkflow, you can then access the values using `workflow.approximator.step_history.history`. If you use an approximator directly, you can find them at `approximator.step_history.history`. This history will not work with the plot function, so you will have to plot the values manually.

---

<div class="post-metadata">

**Author:** ![tdj2000](https://avatars.discourse-cdn.com/v4/letter/t/bbe5ce/32.png) [@tdj2000](https://discuss.bayesflow.org/u/tdj2000)\
**Post date:** [June 7, 2025, 12:57pm UTC](https://discuss.bayesflow.org/t/plotting-of-the-loss/168/3 "2025-06-07T12:57:38Z")

</div>

Thanks for the swift reply and the helpful info!