Keras train_step get epoch number

user_6406072 0
Viewed 1101

I've built my own Keras model class inheriting from keras.Model. However, I have a modified training step, where I manually derive the forward pass and compute two different losses which I sum up such that my total loss is total_loss = loss1 + factor*loss2. I observed that I get better results if I decrease the factor during training. My question would be now, what a good solution for this might be. Can I either access somehow the epoch number in the train_step method such that I can compute the factor from the epoch number or is there any other preferable approach?

thanks a lot :)

EDIT: If I add a counter to the train_step method, it won't increase during training. But this only happens in non-eager mode. If I enable eager mode it works. Any idea on how I could make it work in non-eager mode?

keras tensorflow python
1 Answers

You could be getting better results because you are decreasing the step size as your model approaches a local minima on the loss surface. Notice in the following image how a decreasing step size helps avoid overshooting the minimum. enter image description here

A common approach to implement this is to schedule the learning rate for training ahead of time in a list, where your model accesses the current factor based on the epoch number i.e. factors[epoch]loss2. However, there is nothing stopping you from calculating a factor based directly on the epoch number like factor = numpy.e*(-epoch).

However, your approach of using two losses to create a total loss is not something that I have seen before, there may be some other place in your model where you can directly adjust the learning rate rather than indirectly doing it by adjusting the loss.

user_107428660
Related