on_train_epoch_end how to get current batch??

[ZENGYIMING-EAMON]

on_train_epoch_end how to get the current batch on Cuda device ??

[semaphore-egg]

Are you looking for on_train_batch_end(self, outputs, batch, batch_idx) ?

[ZENGYIMING-EAMON]

if I want to make it inside a epoch end callback ? Is it like:

def on_train_epoch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
print('train outputs: {}'.format(outputs))

?

[ZENGYIMING-EAMON]

I just can’t find the detailed api that where I can get all the info for me to do some predict in callback and i found this based on your answer in another issue.

[ZENGYIMING-EAMON]

Oh I want to do it at epoch end not batch end.

It seems that def on_train_epoch_end(…) can be used inside a lightningmodule or a callback class. Where can I find the detailed aip or instructions to show the difference of these two usages?

[semaphore-egg]

The related hooks are listed here. on_train_batch_end

[ZENGYIMING-EAMON]

how can I predict first batch for every epoch. I found empty args for

def on_train_epoch_end(self) -> None:
"""Called in the training loop at the very end of the epoch.
To access all batch outputs at the end of the epoch, either:
1. Implement training_epoch_end in the LightningModule OR
2. Cache data across steps on the attribute(s) of the LightningModule and access them in this hook
"""

[semaphore-egg]

Not sure about predict first batch, maybe check the batch_idx like:

def on_train_batch_end(self, outputs, batch, batch_idx: int, unused: int = 0) -> None
    if batch_idx == 0:
       # do something for first batch

[ZENGYIMING-EAMON]

Can I call on_train_batch_end inside on_train_epoch_end? because I don’t want model to check after every batch (I want it to check epoch-wise after training all batch)

[semaphore-egg]

The hooks are called automatically by lightning and calling it manually is not encouraged.

Seems you want hooks like:
on_train_first_batch_begin(self, outputs, batch)
on_train_first_batch_end(self, outputs, batch)
on_train_last_batch_begin(self, outputs, batch)
on_train_last_batch_end(self, outputs, batch)

or even
on_train_nth_batch_begin(self, outputs, batch)
on_train_nth_batch_end(self, outputs, batch)

We may put a new issue to request these features in a future version.

[ZENGYIMING-EAMON]

Exactly! Thank you! That would be nice if xxx_epoch_end can utilize n-th batch we want