From ad23361de09df12543912027a42a2cf30ca2f25c Mon Sep 17 00:00:00 2001
From: tomrink <rink@ssec.wisc.edu>
Date: Mon, 25 Oct 2021 14:56:38 -0500
Subject: [PATCH] Use tf.compat.v1.train.Saver to save EWAaveraged trainable
variables
---
modules/deeplearning/icing_cnn.py | 23 +++++++++++------------
1 file changed, 11 insertions(+), 12 deletions(-)
diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 0afe586c..606f881e 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -1,6 +1,6 @@
import tensorflow as tf
import tensorflow_addons as tfa
-from util.setup import logdir, modeldir, cachepath, now
+from util.setup import logdir, modeldir, cachepath, now, ewa_varsdir
from util.util import homedir, EarlyStop, normalize, make_for_full_domain_predict
from util.geos_nav import get_navigation
@@ -72,7 +72,7 @@ train_params_l1b = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'te
# 'cld_emiss_acha', 'conv_cloud_fraction', 'cld_reff_dcomp', 'cld_opd_dcomp', 'iwc_dcomp', 'lwc_dcomp']
# ---------------------------------------------
-train_params = train_params_l1b
+train_params = train_params_l2
# -- Zero out params (Experimentation Only) ------------
zero_out_params = ['cld_reff_dcomp', 'cld_opd_dcomp', 'iwc_dcomp', 'lwc_dcomp']
DO_ZERO_OUT = False
@@ -652,7 +652,7 @@ class IcingIntensityNN:
if TRACK_MOVING_AVERAGE:
# Not really sure this works properly (from tfa)
# optimizer = tfa.optimizers.MovingAverage(optimizer)
- self.ema = tf.train.ExponentialMovingAverage(decay=0.999)
+ self.ema = tf.train.ExponentialMovingAverage(decay=0.9999)
self.optimizer = optimizer
self.initial_learning_rate = initial_learning_rate
@@ -690,9 +690,6 @@ class IcingIntensityNN:
self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
if TRACK_MOVING_AVERAGE:
self.ema.apply(self.model.trainable_variables)
- # TODO: This doesn't seem to work
- # for var in self.model.trainable_variables:
- # var.assign(self.ema.average(var))
self.train_loss(loss)
self.train_accuracy(labels, pred)
@@ -862,12 +859,6 @@ class IcingIntensityNN:
print('loss, acc: ', self.test_loss.result().numpy(), self.test_accuracy.result().numpy())
print('------------------------------------------------------')
- if TRACK_MOVING_AVERAGE:
- # This may not really work properly (from tfa)
- # self.optimizer.assign_average_vars(self.model.trainable_variables)
- for var in self.model.trainable_variables:
- var.assign(self.ema.average(var))
-
tst_loss = self.test_loss.result().numpy()
if tst_loss < best_test_loss:
best_test_loss = tst_loss
@@ -880,6 +871,14 @@ class IcingIntensityNN:
ckpt_manager.save()
+ if TRACK_MOVING_AVERAGE:
+ vars_ema = []
+ for var in self.model.trainable_variables:
+ vars_ema.append(self.ema.average(var))
+
+ saver_ewa = tf.compat.v1.train.Saver(var_list=vars_ema)
+ saver_ewa.save(None, ewa_varsdir)
+
if self.DISK_CACHE and epoch == 0:
f = open(cachepath, 'wb')
pickle.dump(self.in_mem_data_cache, f)
--
GitLab