Merge pull request #26 from emiliocoutinho/tfp_lbfgs

ADAM training on batches, fixed residual implementation for SA-PINN, added testing suite for distributed, non-distributed, adaptive, and baseline PINN examples

Author: levimcclenny

tensordiffeq/fit.py

@@ -14,33 +14,21 @@
 os.environ["TF_GPU_THREAD_MODE"] = "gpu_private"
 
 
-def fit(obj, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
-    # 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:
-        obj.batch_sz = batch_sz
-    else:
-        obj.batch_sz = obj.X_f_len
-        # obj.batch_sz = len(obj.x_f)
-
-    N_f = obj.X_f_len
-    # N_f = len(obj.x_f)
-    n_batches = int(N_f // obj.batch_sz)
+def fit(obj, tf_iter=0, newton_iter=0, newton_eager=True):
+
     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)
 
     # 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)
     if obj.verbose: print_screen(obj)
+
     print("Starting Adam training")
     # tf.profiler.experimental.start('../cache/tblogdir1')
     train_op_fn = train_op_inner(obj)
     with trange(tf_iter) as t:
         for epoch in t:
-            loss_value = train_op_fn(n_batches, obj)
+            loss_value = train_op_fn(obj)
             # Description will be displayed on the left
             t.set_description('Adam epoch %i' % (epoch + 1))
             # Postfix will be displayed on the right,
@@ -53,6 +41,7 @@ def fit(obj, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
 
     # tf.profiler.experimental.start('../cache/tblogdir1')
     if newton_iter > 0:
+        obj.n_batches = 1
         print("Starting L-BFGS training")
         if newton_eager:
             print("Executing eager-mode L-BFGS")
@@ -92,10 +81,14 @@ def lbfgs_op(func, init_params, newton_iter):
 
 
 def train_op_inner(obj):
-    @tf.function()
-    def apply_grads(n_batches, obj=obj):
-        for _ in range(n_batches):
+    @tf.function
+    def apply_grads(obj=obj):
+        if obj.n_batches > 1:
+            obj.batch_indx_map = np.random.choice(obj.X_f_len[0], size=obj.X_f_len[0], replace=False)
+
+        for i in range(obj.n_batches):
             # unstack = tf.unstack(obj.u_model.trainable_variables, axis = 2)
+            obj.batch = i
             obj.variables = obj.u_model.trainable_variables
             obj.variables = obj.u_model.trainable_variables
             if obj.isAdaptive:
@@ -112,7 +105,10 @@ def apply_grads(n_batches, obj=obj):
             else:
                 loss_value, grads = obj.grad()
                 obj.tf_optimizer.apply_gradients(zip(grads, obj.u_model.trainable_variables))
-            return loss_value
+
+        obj.batch = None
+
+        return loss_value
 
     return apply_grads
 

tensordiffeq/models.py

@@ -48,6 +48,8 @@ def compile(self, layer_sizes, f_model, domain, bcs, isAdaptive=False,
         # self.X_f_in = np.asarray(tmp)
         self.X_f_in = [tf.cast(np.reshape(vec, (-1, 1)), tf.float32) for i, vec in enumerate(self.domain.X_f.T)]
         self.u_model = neural_net(self.layer_sizes)
+        self.batch = None
+        self.batch_indx_map = None
         self.lambdas = self.dict_adaptive = self.lambdas_map = None
         self.isAdaptive = isAdaptive
 
@@ -83,11 +85,12 @@ def update_loss(self):
         #####################################
         # Check if adaptive is allowed
         if self.isAdaptive:
-            idx_lambda_bcs = self.lambdas_map['bcs'][0]
+            if len(self.lambdas_map['bcs']) > 0:
+                idx_lambda_bcs = self.lambdas_map['bcs'][0]
 
         for counter_bc, bc in enumerate(self.bcs):
             loss_bc = 0.
-            # Check if the current BS is adaptive
+            # Check if the current BC is adaptive
             if self.isAdaptive:
                 isBC_adaptive = self.dict_adaptive["BCs"][counter_bc]
             else:
@@ -142,7 +145,16 @@ def update_loss(self):
         # Residual Equations
         #####################################
         # pass thorough the forward method
-        f_u_preds = self.f_model(self.u_model, *self.X_f_in)
+        if self.n_batches > 1:
+            # The collocation points will be split based on the batch_indx_map
+            # generated on the beginning of this epoch on models.train_op_inner.apply_grads
+            X_batch = []
+            for x_in in self.X_f_in:
+                indx_on_batch = self.batch_indx_map[self.batch * self.batch_sz:(self.batch + 1) * self.batch_sz]
+                X_batch.append(tf.gather(x_in,indx_on_batch))
+            f_u_preds = self.f_model(self.u_model, *X_batch)
+        else:
+            f_u_preds = self.f_model(self.u_model, *self.X_f_in)
 
         # If it is only one residual, just convert it to a tuple of one element
         if not isinstance(f_u_preds, tuple):
@@ -153,13 +165,23 @@ def update_loss(self):
             # Check if the current Residual is adaptive
             if self.isAdaptive:
                 isRes_adaptive = self.dict_adaptive["residual"][counter_res]
-                idx_lambda_res = self.lambdas_map['residual'][0]
                 if isRes_adaptive:
+                    idx_lambda_res = self.lambdas_map['residual'][0]
+                    lambdas2loss = self.lambdas[idx_lambda_res]
+
+                    if self.n_batches > 1:
+                        # select lambdas on minebatch
+                        lambdas2loss = tf.gather(lambdas2loss,indx_on_batch)
+
                     if self.g is not None:
-                        loss_r = g_MSE(f_u_pred, constant(0.0), self.g(self.lambdas[idx_lambda_res]))
+                        loss_r = g_MSE(f_u_pred, constant(0.0), self.g(lambdas2loss))
                     else:
-                        loss_r = MSE(f_u_pred, constant(0.0), self.lambdas[idx_lambda_res])
+                        loss_r = MSE(f_u_pred, constant(0.0), lambdas2loss)
                     idx_lambda_res += 1
+                else:
+                    # In the case where the model is Adaptive but the residual
+                    # is not adaptive, the residual loss should be computed.
+                    loss_r = MSE(f_u_pred, constant(0.0))
             else:
                 loss_r = MSE(f_u_pred, constant(0.0))
 
@@ -177,8 +199,18 @@ def grad(self):
         return loss_value, grads
 
     def fit(self, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
-        if self.isAdaptive and (batch_sz is not None):
-            raise Exception("Currently we dont support minibatching for adaptive PINNs")
+
+        # Can adjust batch size for collocation points, here we set it to N_f
+        N_f = self.X_f_len[0]
+        self.batch_sz = batch_sz if batch_sz is not None else N_f
+        self.n_batches = N_f // self.batch_sz
+
+        if self.isAdaptive and self.dist:
+            raise Exception("Currently we dont support distributed training for adaptive PINNs")
+
+        if self.n_batches > 1 and self.dist:
+            raise Exception("Currently we dont support distributed minibatching training")
+
         if self.dist:
             BUFFER_SIZE = len(self.X_f_in[0])
             EPOCHS = tf_iter
@@ -194,13 +226,6 @@ def fit(self, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
 
             print("Number of GPU devices: {}".format(self.strategy.num_replicas_in_sync))
 
-            self.batch_sz = batch_sz if batch_sz is not None else len(self.X_f_in[0])
-            # weights_idx = tensor(list(range(len(self.x_f))), dtype=tf.int32)
-            # print(weights_idx)
-            # print(tf.gather(self.col_weights, weights_idx))
-            N_f = len(self.X_f_in[0])
-            self.n_batches = N_f // self.batch_sz
-
             BATCH_SIZE_PER_REPLICA = self.batch_sz
             GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * self.strategy.num_replicas_in_sync
 
@@ -229,7 +254,7 @@ def fit(self, tf_iter=0, newton_iter=0, batch_sz=None, newton_eager=True):
             fit_dist(self, tf_iter=tf_iter, newton_iter=newton_iter, batch_sz=batch_sz, newton_eager=newton_eager)
 
         else:
-            fit(self, tf_iter=tf_iter, newton_iter=newton_iter, batch_sz=batch_sz, newton_eager=newton_eager)
+            fit(self, tf_iter=tf_iter, newton_iter=newton_iter, newton_eager=newton_eager)
 
     # L-BFGS implementation from https://github.com/pierremtb/PINNs-TF2.0
     def get_loss_and_flat_grad(self):

tensordiffeq/test/AC2test.py

@@ -0,0 +1,117 @@
+import pytest
+from tensordiffeq.boundaries import *
+import scipy.io
+import math
+import tensordiffeq as tdq
+from tensordiffeq.models import CollocationSolverND
+
+def main(args):
+
+    if args is None:
+        args = {'layer_sizes': [2, 21, 21, 21, 21, 1],
+
+                'run_functions_eagerly': False,
+                'epoch_adam': 20,
+                'epoch_lbfgs': 20,
+                'lbfgs_eager': False,
+                'isAdaptive': True,
+                'dist_training': False,
+                'dict_adaptive': {"residual": [True],
+                                  "BCs": [False, False]},
+                'N_x': 100,
+                'N_t': 50,
+                'N_f': 5000,
+                'batch_sz': 200,
+                }
+
+    layer_sizes = args['layer_sizes']
+    run_functions_eagerly = args['run_functions_eagerly']
+    epoch_adam = args['epoch_adam']
+    epoch_lbfgs = args['epoch_lbfgs']
+    lbfgs_eager = args['lbfgs_eager']
+    isAdaptive = args['isAdaptive']
+    dist_training = args['dist_training']
+    dict_adaptive = args['dict_adaptive']
+    N_x = args['N_x']
+    N_t = args['N_t']
+    N_f = args['N_f']
+    batch_sz = args['batch_sz']
+
+
+    tf.config.run_functions_eagerly(run_functions_eagerly)
+
+    Domain = DomainND(["x", "t"], time_var='t')
+
+    Domain.add("x", [-1.0, 1.0], N_x)
+    Domain.add("t", [0.0, 1.0], N_t)
+
+    Domain.generate_collocation_points(N_f)
+
+
+    def func_ic(x):
+        return x ** 2 * np.cos(math.pi * x)
+
+
+    # Conditions to be considered at the boundaries for the periodic BC
+    def deriv_model(u_model, x, t):
+        u = u_model(tf.concat([x, t], 1))
+        u_x = tf.gradients(u, x)[0]
+
+        return u, u_x
+
+
+    init = IC(Domain, [func_ic], var=[['x']])
+    x_periodic = periodicBC(Domain, ['x'], [deriv_model])
+
+    BCs = [init, x_periodic]
+
+
+    def f_model(u_model, x, t):
+        u = u_model(tf.concat([x, t], 1))
+        u_x = tf.gradients(u, x)
+        u_xx = tf.gradients(u_x, x)
+        u_t = tf.gradients(u, t)
+        c1 = tdq.utils.constant(.0001)
+        c2 = tdq.utils.constant(5.0)
+        f_u = u_t - c1 * u_xx + c2 * u * u * u - c2 * u
+        return f_u
+
+    ## Which loss functions will have adaptive weights
+    # "residual" should a tuple for the case of multiple residual equation
+    # BCs have to follow the same order as the previously defined BCs list
+    dict_adaptive = dict_adaptive
+
+    ## Weights initialization
+    # dictionary with keys "residual" and "BCs". Values must be a tuple with dimension
+    # equal to the number of residuals and boundary conditions, respectively
+
+    if dict_adaptive["residual"][0] == False:
+        init_residual = None
+    else:
+        init_residual = tf.random.uniform([N_f, 1])
+
+    if dict_adaptive["BCs"][0] == False:
+        init_IC = None
+    else:
+        init_IC = 100 * tf.random.uniform([N_x, 1])
+
+    if dict_adaptive["BCs"][1] == False:
+        init_BC = None
+    else:
+        init_BC = tf.random.uniform([N_t, 1])
+
+    init_weights = {"residual": [init_residual],
+                    "BCs": [init_IC, init_BC]}
+
+
+
+    model = CollocationSolverND()
+    model.compile(layer_sizes, f_model, Domain, BCs, isAdaptive=isAdaptive,
+                  dict_adaptive=dict_adaptive, init_weights=init_weights, dist=dist_training)
+
+    model.fit(tf_iter=epoch_adam, newton_iter=epoch_lbfgs, newton_eager=lbfgs_eager, batch_sz=batch_sz)
+
+    return
+
+if __name__ == "__main__":
+    main(args=None)