Merge pull request #18 from levimcclenny/main

Fix output of discoverymodel

Author: levimcclenny

tensordiffeq/fit.py

@@ -42,19 +42,14 @@ def fit(obj, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
         for epoch in t:
             loss_value = train_op_fn(n_batches, obj)
             # Description will be displayed on the left
-            t.set_description('Adam epoch %i' % (epoch+1))
+            t.set_description('Adam epoch %i' % (epoch + 1))
             # Postfix will be displayed on the right,
             # formatted automatically based on argument's datatype
             if epoch % 10 == 0:
-                elapsed = time.time() - start_time
                 t.set_postfix(loss=loss_value.numpy())
-                #
-                # print('It: %d, Time: %.2f, loss: %.6f' % (epoch, elapsed, loss_value.numpy()))
-                # tf.print(f"total loss: {loss_value}")
-                start_time = time.time()
 
-    # tf.profiler.experimental.stop()
 
+    # tf.profiler.experimental.stop()
 
     # tf.profiler.experimental.start('../cache/tblogdir1')
     if newton_iter > 0:
@@ -97,7 +92,7 @@ def lbfgs_op(func, init_params, newton_iter):
 
 
 def train_op_inner(obj):
-    @tf.function
+    @tf.function()
     def apply_grads(n_batches, obj=obj):
         for _ in range(n_batches):
             # unstack = tf.unstack(obj.u_model.trainable_variables, axis = 2)
@@ -126,7 +121,6 @@ def apply_grads(n_batches, obj=obj):
     return apply_grads
 
 
-
 def fit_dist(obj, tf_iter, newton_iter, batch_sz=None, newton_eager=True):
     def train_epoch(dataset, STEPS):
         total_loss = 0.0
@@ -180,17 +174,17 @@ def dist_loop(obj, STEPS):
         return train_loss
 
     def train_loop(obj, tf_iter, STEPS):
-
+        print_screen(obj)
         start_time = time.time()
-        for epoch in range(tf_iter):
-            loss = dist_loop(obj, STEPS)
-
-            if epoch % 100 == 0:
-                elapsed = time.time() - start_time
-                template = ("Epoch {}, Time: {}, Loss: {}")
-                print(template.format(epoch, elapsed, loss))
-                # print('It: %d, Time: %.2f, loss: %.9f' % (epoch, elapsed, tf.get_static_value(loss)))
-                start_time = time.time()
+        with trange(tf_iter) as t:
+            for epoch in t:
+                loss = dist_loop(obj, STEPS)
+                t.set_description('Adam epoch %i' % (epoch + 1))
+                if epoch % 10 == 0:
+                    elapsed = time.time() - start_time
+                    t.set_postfix(loss=loss.numpy())
+                    # print('It: %d, Time: %.2f, loss: %.9f' % (epoch, elapsed, tf.get_static_value(loss)))
+                    start_time = time.time()
 
     print("starting Adam training")
     STEPS = np.max((obj.n_batches // obj.strategy.num_replicas_in_sync, 1))

tensordiffeq/models.py

@@ -5,6 +5,8 @@
 from .networks import *
 from .plotting import *
 from .fit import *
+from tqdm.auto import tqdm, trange
+from .output import print_screen
 
 
 class CollocationSolverND:
@@ -232,10 +234,7 @@ def compile(self, layer_sizes, f_model, X, u, var, col_weights=None):
     @tf.function
     def loss(self):
         u_pred = self.u_model(tf.concat(self.X, 1))
-        print(self.vars)
         f_u_pred = self.f_model(self.u_model, self.vars, *self.X_in)
-        print(self.vars)
-
         if self.col_weights is not None:
             return MSE(u_pred, self.u) + g_MSE(f_u_pred, constant(0.0), self.col_weights ** 2)
         else:
@@ -273,15 +272,18 @@ def train_op(self):
     def fit(self, tf_iter):
         self.train_loop(tf_iter)
 
-    # @tf.function
     def train_loop(self, tf_iter):  # sourcery skip: move-assign
         start_time = time.time()
-        for i in range(tf_iter):
-            loss_value = self.train_op()
-            if i % 100 == 0:
-                elapsed = time.time() - start_time
-                print('It: %d, Time: %.2f' % (i, elapsed))
-                tf.print(f"loss_value: {loss_value}")
-                var = [var.numpy() for var in self.vars]
-                tf.print(f"vars estimate(s): {var}")
-                start_time = time.time()
+        print_screen(self, discovery_model=True)
+        with trange(tf_iter) as t:
+            for i in t:
+                loss_value = self.train_op()
+                if i % 10 == 0:
+                    # elapsed = time.time() - start_time
+                    # print('It: %d, Time: %.2f' % (i, elapsed))
+                    # tf.print(f"loss_value: {loss_value}")
+                    var = [var.numpy() for var in self.vars]
+                    t.set_postfix(loss=loss_value.numpy())
+                    t.set_postfix(vars=var)
+                    # tf.print(f"vars estimate(s): {var}")
+                    # start_time = time.time()

tensordiffeq/output.py

@@ -2,8 +2,10 @@
 from os import system, name
 import sys
 
-def print_screen(model):
+def print_screen(model, discovery_model=False):
     f = Figlet(font='slant')
     print(f.renderText('TensorDiffEq'))
+    if discovery_model:
+        print("Running Discovery Model for Parameter Estimation\n\n")
     print("Neural Network Model Summary\n")
     print(model.u_model.summary())