diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py index fc0ad7d87113f817f1ac66fba96ccbe14660cf06..f7c464e44b8bc6d90f3a82c68a3be0f2f0f5712f 100644 --- a/modules/deeplearning/icing_cnn.py +++ b/modules/deeplearning/icing_cnn.py @@ -1,4 +1,5 @@ import tensorflow as tf +import tensorflow_addons as tfa from util.setup import logdir, modeldir, cachepath from util.util import homedir import subprocess @@ -18,7 +19,7 @@ PROC_BATCH_SIZE = 2046 PROC_BATCH_BUFFER_SIZE = 50000 NumLabels = 1 BATCH_SIZE = 256 -NUM_EPOCHS = 100 +NUM_EPOCHS = 60 TRACK_MOVING_AVERAGE = False @@ -400,6 +401,7 @@ class IcingIntensityNN: self.learningRateSchedule = tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate, decay_steps, decay_rate) optimizer = tf.keras.optimizers.Adam(learning_rate=self.learningRateSchedule) + optimizer = tfa.optimizers.MovingAverage(optimizer) if TRACK_MOVING_AVERAGE: ema = tf.train.ExponentialMovingAverage(decay=0.999)