tfa.callbacks.AverageModelCheckpoint

View source on GitHub

The callback that should be used with optimizers that extend

AverageWrapper, i.e., MovingAverage and StochasticAverage optimizers. It saves and, optionally, assigns the averaged weights.

Args:

  • update_weights: If True, assign the moving average weights to the model, and save them. If False, keep the old non-averaged weights, but the saved model uses the average weights.

See tf.keras.callbacks.ModelCheckpoint for the other args.

Methods

set_model

View source

set_params