decouple model.compile and model.fit to allow repeated calling of fit

Author: levimcclenny

.gitignore

@@ -2,4 +2,5 @@ cache/*
 tensordiffeq/__pycache__/*
 venv/*
 .idea/*
-dist/*
+dist/*
+examples/*.ipynb

examples/burgers-new.py

@@ -38,7 +38,9 @@ def f_model(u_model, x, t):
 
 model = CollocationSolverND()
 model.compile(layer_sizes, f_model, Domain, BCs)
-model.fit(tf_iter=1000, newton_iter=1000)
+model.fit(tf_iter=301, newton_iter=101)
+
+model.fit(tf_iter=301, newton_iter=101)
 
 
 #######################################################

tensordiffeq/fit.py

@@ -12,7 +12,7 @@
 
 
 def fit(obj, tf_iter, newton_iter, batch_sz=None, newton_eager=True):
-    obj.u_model = neural_net(obj.layer_sizes)
+    # obj.u_model = neural_net(obj.layer_sizes)
     # obj.build_loss()
     # Can adjust batch size for collocation points, here we set it to N_f
     if batch_sz is not None:
@@ -25,19 +25,18 @@ def fit(obj, tf_iter, newton_iter, batch_sz=None, newton_eager=True):
     # N_f = len(obj.x_f)
     n_batches = int(N_f // obj.batch_sz)
     start_time = time.time()
-    obj.tf_optimizer = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
-    obj.tf_optimizer_weights = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
+    # obj.tf_optimizer = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
+    # obj.tf_optimizer_weights = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
 
     # these cant be tf.functions on initialization since the distributed strategy requires its own
     # graph using grad and adaptgrad, so they cant be compiled as tf.functions until we know dist/non-dist
     obj.grad = tf.function(obj.grad)
     print("starting Adam training")
     # tf.profiler.experimental.start('../cache/tblogdir1')
-    print(n_batches)
-    print(tf_iter)
+    train_op_fn = train_op_inner(obj)
+    print(obj.tf_optimizer)
     for epoch in range(tf_iter):
-
-        loss_value = train_op(obj, n_batches)
+        loss_value = train_op_fn(n_batches, obj)
 
         if epoch % 100 == 0:
             elapsed = time.time() - start_time
@@ -85,23 +84,29 @@ def lbfgs_op(func, init_params, newton_iter):
         tolerance=1e-20,
     )
 
-@tf.function()
-def train_op(obj, n_batches):
-    for _ in range(n_batches):
-        # unstack = tf.unstack(obj.u_model.trainable_variables, axis = 2)
-        obj.variables = obj.u_model.trainable_variables
-        if obj.isAdaptive:
-            obj.variables.extend([obj.u_weights, obj.col_weights])
-            loss_value, grads = obj.grad()
-            obj.tf_optimizer.apply_gradients(zip(grads[:-2], obj.u_model.trainable_variables))
-            obj.tf_optimizer_weights.apply_gradients(zip([-grads[-2], -grads[-1]], [obj.u_weights, obj.col_weights]))
-        else:
-            loss_value, grads = obj.grad()
-            obj.tf_optimizer.apply_gradients(zip(grads, obj.u_model.trainable_variables))
-        return loss_value
+
+def train_op_inner(obj):
+    @tf.function
+    def apply_grads(n_batches, obj=obj):
+        for _ in range(n_batches):
+            # unstack = tf.unstack(obj.u_model.trainable_variables, axis = 2)
+            obj.variables = obj.u_model.trainable_variables
+            if obj.isAdaptive:
+                obj.variables.extend([obj.u_weights, obj.col_weights])
+                loss_value, grads = obj.grad()
+                obj.tf_optimizer.apply_gradients(zip(grads[:-2], obj.u_model.trainable_variables))
+                obj.tf_optimizer_weights.apply_gradients(
+                    zip([-grads[-2], -grads[-1]], [obj.u_weights, obj.col_weights]))
+            else:
+                loss_value, grads = obj.grad()
+                obj.tf_optimizer.apply_gradients(zip(grads, obj.u_model.trainable_variables))
+            return loss_value
+
+    return apply_grads
 
 
 # TODO Distributed training re-integration
+# TODO decouple u_model from being overwritten by calling model.fit
 
 def fit_dist(obj, tf_iter, newton_iter, batch_sz=None, newton_eager=True):
     BUFFER_SIZE = len(obj.x_f)

tensordiffeq/models.py

@@ -13,6 +13,8 @@ def __init__(self, assimilate=False):
 
     def compile(self, layer_sizes, f_model, domain, bcs, isAdaptive=False,
                 col_weights=None, u_weights=None, g=None, dist=False):
+        self.tf_optimizer = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
+        self.tf_optimizer_weights = tf.keras.optimizers.Adam(lr=0.005, beta_1=.99)
         self.layer_sizes = layer_sizes
         self.sizes_w, self.sizes_b = get_sizes(layer_sizes)
         self.bcs = bcs
@@ -27,6 +29,7 @@ def compile(self, layer_sizes, f_model, domain, bcs, isAdaptive=False,
         self.X_f_len = tf.slice(self.X_f_dims, [0], [1]).numpy()
         tmp = [np.reshape(vec, (-1,1)) for i, vec in enumerate(self.domain.X_f.T)]
         self.X_f_in = np.asarray(tmp)
+        self.u_model = neural_net(self.layer_sizes)
 
 
 
@@ -87,6 +90,7 @@ def update_loss(self):
         loss_tmp = tf.math.add(loss_tmp, mse_f_u)
         return loss_tmp
 
+    #@tf.function
     def grad(self):
         with tf.GradientTape() as tape:
             loss_value = self.update_loss()