minor...

modules/deeplearning/icing_cnn.py

@@ -21,7 +21,7 @@ NumLogits = 1
 BATCH_SIZE = 256
 NUM_EPOCHS = 200
 
-TRACK_MOVING_AVERAGE = False
+TRACK_MOVING_AVERAGE = True
 
 TRIPLET = False
 CONV3D = False
@@ -462,13 +462,9 @@ 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)
-
-            with tf.control_dependencies([optimizer]):
-                optimizer = ema.apply(self.model.trainable_variables)
+            optimizer = tfa.optimizers.MovingAverage(optimizer)
 
         self.optimizer = optimizer
         self.initial_learning_rate = initial_learning_rate
@@ -498,14 +494,6 @@ class IcingIntensityNN:
             self.test_false_neg = tf.keras.metrics.FalseNegatives(name='test_false_neg')
             self.test_false_pos = tf.keras.metrics.FalsePositives(name='test_false_pos')
 
-    def build_predict(self):
-        _, pred = tf.nn.top_k(self.logits)
-        self.pred_class = pred
-
-        if TRACK_MOVING_AVERAGE:
-            self.variable_averages = tf.train.ExponentialMovingAverage(0.999, self.global_step)
-            self.variable_averages.apply(self.model.trainable_variables)
-
     @tf.function
     def train_step(self, mini_batch):
         inputs = [mini_batch[0]]