Accessing the current training iteration step in a custom Layer class

[mpetteno]

Hi everyone,

Is there a way to access to the current training iteration step when using a custom Layer class? Currently the only way I found is to pass it from the Model class when calling the layer, like this:

class TestModel(keras.Model):

    def __init__(self, name: str = "test_model", **kwargs):
        super(VAE, self).__init__(name=name, **kwargs)
        self._custom_layer = MyLayer()

    def call(self, inputs, training: bool = False):
        current_step = self.optimizer.iterations + 1
        return self._custom_layer(training=training, current_step=current_step)

This is ok, but in my opinion it is a little bit redundant if the custom layer defines other layers that need the parameter as well. It will be convenient to have the optimizer available at layer-level for example. I think that this will be useful in multi-objective scenarios for example, where often you want to anneal the hyperparameter that multiplies a term in the total loss function.

[fchollet]

Currently the only way I found is to pass it from the Model class when calling the layer, like this:

That is fine.

Another thing you can do is make the layer keep its own local iteration counter, incremented every time the layer is called in training mode.

If you want access to the optimizer, you can always create the optimizer before model construction and pass it to your custom layer. Then call compile() with the same optimizer instance.

[mpetteno]

Another thing you can do is make the layer keep its own local iteration counter, incremented every time the layer is called in training mode.

What If the layers are called more than once in a single training iteration?

Anyway my point is that it is fine if you only have to pass the counter to a single layer, but It can be cumbersome if there are nested layers. But I guess that it is not trivial to achive this given the actual relations between Model, Trainer and Layer right?

[fchollet]

Optimizers live at the model level. For a layer to be aware of the optimizer, the optimizer must be provided to the layer manually. Or else, place the logic that needs to be optimizer-aware at the model level (e.g. in Model.call).

[mpetteno]

Ok, thanks.