Calculating a loss function over the whole training set and updating the model

[Sadegh-Saberian]

Hello,

I have a loss function that should be calculated over the whole training set instead of batches. I was trying to implement it using training_epoch_end, but realized that this function is mostly for logging purposes. My question is how can I implement this loss so that it is involved in optimization?

[rohitgr7]

I was trying to implement it using training_epoch_end, but realized that this function is mostly for logging purposes

not necessary, you can do anything in that. If you want to accumulate the loss for each batch over the complete epoch and log it, you can simply use:

def training_step(self, batch, batch_idx):
    loss = ...
    self.log('train_loss', loss, on_step=False, on_epoch=True)
    return loss

or in case you want to do it manually by yourself, you can return it inside training_step and accumulate it inside training_epoch_end.

[tchaton]

Hey @Sadegh-Saberian

should be calculated over the whole training set instead of batches

Do you mean you can your loss be computed at once over the entire dataset? Would your entire dataset could fit as a single batch?

An alternative solution would be to do Trainer(accumulate_grad_batches=len(train_dataset) / batch_size). So your gradients get accumulated over your entire dataset before updating the weights, therefore 1 epoch = 1 model update.

[Sadegh-Saberian]

Thanks for your replies. I think accumulating the grad over batches is the way to go! thanks!