Applying Levenberg-Marquardt to physically informed neural networks (PINNs)

[JohnTaylor2000]

Applying Levenberg-Marquardt to physically informed neural networks (PINNs)

[fabiodimarco]

Hi John,
very interesting I'll read the paper. As regards the PINNs network could you provide me a zip with the code so that I can easily replicate the error?
You can send it to my email [email protected]

[JohnTaylor2000]

In the code https://github.com/okada39/pinn_burgers in main.py I added the following code to replace the call to lbfgs. The call to lm.ModelWrapper generates the error:-

lbfgs = L_BFGS_B(model=pinn, x_train=x_train, y_train=y_train)

lbfgs.fit()

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
#train_dataset = train_dataset.shuffle(input_size)
train_dataset = train_dataset.batch(batch_size).cache()
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

model_wrap = lm.ModelWrapper(pinn)
model_wrap.compile( optimizer=tf.keras.optimizers.SGD(learning_rate=1.0),      loss=lm.MeanSquaredError()
        ,experimental_run_tf_function=False)
print("\n_________________________________________________________________")
print("Train using Levenberg-Marquardt")
t1_start = time.perf_counter()
history = model_wrap.fit(x_train, y_train, epochs=150, verbose=2)
t1_stop = time.perf_counter()
print("Levenberg-Marquardt Elapsed time: ", t1_stop - t1_start)
eval    = model_wrap.evaluate(train_dataset, verbose =2)

[fabiodimarco]

The problem is due to the model having inputs and outputs defined as lists. At the moment they can only be tensors, because the ModelWrapper class inherits from tf.keras.Sequential (I'll have to solve this in the future).
However, you can easly overcome the problem by splitting and concatenating your data inside the model definition:

        x = tf.keras.layers.Input(shape=(6,))

        [tx_eqn, tx_ini, tx_bnd] = tf.keras.layers.Lambda(
            lambda z: tf.split(z, 3, axis=-1))(x)

        # compute gradients
        u, du_dt, du_dx, d2u_dx2 = self.grads(tx_eqn)

        # equation output being zero
        u_eqn = du_dt + u*du_dx - self.nu*d2u_dx2
        # initial condition output
        u_ini = self.network(tx_ini)
        # boundary condition output
        u_bnd = self.network(tx_bnd)

        y = tf.keras.layers.Lambda(
            lambda z: tf.concat(z, axis=-1))([u_eqn, u_ini, u_bnd])

        # build the PINN model for Burgers' equation
        return tf.keras.models.Model(inputs=x, outputs=y)

Let me know if this solves your problem.

[JohnTaylor2000]

That solves the problem with the wrapper!

Unfortunately, a problem has now popped up in the fit routine:-

Epoch 1/10
Traceback (most recent call last):
File "/home/599/jt9268/PINN_project/pinn_burgers/main.py", line 72, in
history = model_wrap.fit(train_dataset, epochs=10, verbose=2)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 746, in fit
return super(ModelWrapper, self).fit(
File "/scratch/ue12/jt9268/tflow/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 713, in train_step
self.trainer.train_step(inputs, targets)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 568, in train_step
self._num_outputs = self._compute_num_outputs(inputs, targets)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 549, in _compute_num_outputs
outputs = self.model(_inputs)
File "/home/599/jt9268/PINN_project/pinn_burgers/lib/pinn.py", line 43, in
lambda z: tf.split(z, 3, axis=-1))(x)
ValueError: Exception encountered when calling layer "lambda" (type Lambda).

Dimension size must be evenly divisible by 3 but is 2
Number of ways to split should evenly divide the split dimension for '{{node model_1/lambda/split}} = Split[T=DT_FLOAT, num_split=3](model_1/lambda/split/split_dim, model_1/Cast)' with input shapes: [], [?,12000,2] and with computed input tensors: input[0] = <-1>.

Call arguments received:
• inputs=tf.Tensor(shape=(None, 12000, 2), dtype=float32)
• mask=None
• training=None

[fabiodimarco]

Have you adjusted the dimensions of the input and output data according to the new model definition?

x_train = tf.concat([tx_eqn, tx_ini, tx_bnd], axis=-1)
y_train = tf.concat([u_eqn,  u_ini,  u_bnd], axis=-1)

[JohnTaylor2000]

Success - I have the code running and producing the correct solution - I have attached a plot below! Many thanks for your advice. I would be interested to get your comments on the paper if you have time?

I did encounter an additional error (see below) and needed change the code to:-

x_train = tf.cast(tf.concat([tx_eqn, tx_ini, tx_bnd], axis=-1),tf.float32)
y_train = tf.cast(tf.concat([u_eqn, u_ini, u_bnd], axis=-1),tf.float32)

Traceback (most recent call last):
File "/home/599/jt9268/PINN_project/pinn_burgers/main.py", line 76, in
history = model_wrap.fit(x_train, y_train, epochs=10, verbose=2)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 746, in fit
return super(ModelWrapper, self).fit(
File "/scratch/ue12/jt9268/tflow/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 713, in train_step
self.trainer.train_step(inputs, targets)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 568, in train_step
self._num_outputs = self._compute_num_outputs(inputs, targets)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 550, in _compute_num_outputs
residuals = self.loss.residuals(_targets, outputs)
File "/home/599/jt9268/PINN_project/pinn_burgers/levenberg_marquardt.py", line 35, in residuals
return y_true - y_pred
TypeError: Exception encountered when calling layer "tf.math.subtract_1" (type TFOpLambda).

Input 'y' of 'Sub' Op has type float32 that does not match type float64 of argument 'x'.
plot_pinn.pdf

[fabiodimarco]

In regards to this part of the paper:

In terms of implementing LM on GPU (e.g. V100), we are able to get the LM algorithm to train with
1024 hidden units (over 1 million parameters) with GPU utilization rising to 11%. At 2048 hidden units
(4.2 million parameters) GPU utilization rose to 35%. At 4096 hidden units (16 million parameters)
GPU utilization rose to 70%. It topped out around 83% with 6000 hidden units (36 million parameters). 

I assume you used LM in a small batch size configuration with num_residuals < num_weights (underdetermined).
In my experiments this configuration did not provide great results compared to the overdetermined case.
What was the batch size? Did you use fixed damping factor or adaptive damping factor?

[JohnTaylor2000]

For these experiments we were just interested in GPU utilisation. On the small networks that we use for anomaly detection , which is partly a consequence of using Lavenberg-Marquardt, GPU utilisation is low, and there is not much we can do about this. I was demonstrating just how big a problem is required to fill a GPU. I did not look at the quality of the solution as we would not use these large models for anomaly detection. Thanks for taking a look at the paper.

[fabiodimarco]

Okok, let me know if you find other intereting application to LM.
Given that the issue is solved I'll close it then.