lululxvi / deeponet

Learning nonlinear operators via DeepONet
Other
509 stars 139 forks source link

Tests not running properly #11

Closed bhaveshshrimali closed 2 years ago

bhaveshshrimali commented 3 years ago

Hi @lululxvi , Thanks for putting together deepxde. I was trying to get started with ODESystem but noticed a couple of things that are off. First off, here is my version info:

In [1]: import tensorflow, deepxde

In [2]: tensorflow.__version__
Out[2]: '2.6.0'

In [3]: deepxde.__version__
Out[3]: '0.13.5'

Now if I try to run the gravity pendulum ODE, I get a weird type mismatch error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-359904db8a4e> in <module>()
     47 
     48 if __name__ == "__main__":
---> 49     main()

14 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    992           except Exception as e:  # pylint:disable=broad-except
    993             if hasattr(e, "ag_error_metadata"):
--> 994               raise e.ag_error_metadata.to_exception(e)
    995             else:
    996               raise

TypeError: in user code:

    /usr/local/lib/python3.7/dist-packages/deepxde/model.py:169 outputs_losses  *
        losses = self.data.losses(targets, outputs_, loss_fn, self)
    /usr/local/lib/python3.7/dist-packages/deepxde/data/triple.py:32 losses  *
        return [loss(targets, outputs)]
    /usr/local/lib/python3.7/dist-packages/deepxde/losses.py:27 mean_squared_error  *
        return bkd.reduce_mean(bkd.square(y_true - y_pred))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py:1383 binary_op_wrapper
        raise e
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py:1367 binary_op_wrapper
        return func(x, y, name=name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:206 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py:548 subtract
        return gen_math_ops.sub(x, y, name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/gen_math_ops.py:10654 sub
        "Sub", x=x, y=y, name=name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/op_def_library.py:558 _apply_op_helper
        inferred_from[input_arg.type_attr]))

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

Here is the piece of code that I am trying to run (I cloned and manually appended the path to the git repo to directly import some of the submodules)

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools

import numpy as np
import sys, os

sys.path.append("deeponet/src/")

import deepxde as dde
from spaces import FinitePowerSeries, FiniteChebyshev, GRF
from system import LTSystem, ODESystem, DRSystem, CVCSystem, ADVDSystem
from utils import merge_values, trim_to_65535, mean_squared_error_outlier, safe_test

def test_u_ode(nn, system, T, m, model, data, u, fname, num=100):
    """Test ODE"""
    sensors = np.linspace(0, T, num=m)[:, None]
    sensor_values = u(sensors)
    x = np.linspace(0, T, num=num)[:, None]
    X_test = [np.tile(sensor_values.T, (num, 1)), x]
    y_test = system.eval_s_func(u, x)
    if nn != "opnn":
        X_test = merge_values(X_test)
    y_pred = model.predict(data.transform_inputs(X_test))
    np.savetxt(fname, np.hstack((x, y_test, y_pred)))
    print("L2relative error:", dde.metrics.l2_relative_error(y_test, y_pred))

def ode_system(T):
    """ODE"""

    def g(s, u, x):
        # Gravity pendulum
        k = 1
        return [s[1], -k * np.sin(s[0]) + u]

    s0 = [0, 0]  # Gravity pendulum
    return ODESystem(g, s0, T)

def run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test):
    # space_test = GRF(1, length_scale=0.1, N=1000, interp="cubic")

    X_train, y_train = system.gen_operator_data(space, m, num_train)
    X_test, y_test = system.gen_operator_data(space, m, num_test)
    if nn != "opnn":
        X_train = merge_values(X_train)
        X_test = merge_values(X_test)

    # np.savez_compressed("train.npz", X_train0=X_train[0], X_train1=X_train[1], y_train=y_train)
    # np.savez_compressed("test.npz", X_test0=X_test[0], X_test1=X_test[1], y_test=y_test)
    # return

    # d = np.load("train.npz")
    # X_train, y_train = (d["X_train0"], d["X_train1"]), d["y_train"]
    # d = np.load("test.npz")
    # X_test, y_test = (d["X_test0"], d["X_test1"]), d["y_test"]

    X_test_trim = trim_to_65535(X_test)[0]
    y_test_trim = trim_to_65535(y_test)[0]
    if nn == "opnn":
        data = dde.data.Triple(
            X_train=X_train, y_train=y_train, X_test=X_test_trim, y_test=y_test_trim
        )
    else:
        data = dde.data.DataSet(
            X_train=X_train, y_train=y_train, X_test=X_test_trim, y_test=y_test_trim
        )

    model = dde.Model(data, net)
    model.compile("adam", lr=lr, metrics=[mean_squared_error_outlier])
    checker = dde.callbacks.ModelCheckpoint(
        "model/model.ckpt", save_better_only=True, period=1000
    )
    losshistory, train_state = model.train(epochs=epochs, callbacks=[checker])
    print(
        "# Parameters:",
        np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]),
    )
    dde.saveplot(losshistory, train_state, issave=True, isplot=True)

    model.restore("model/model.ckpt-" + str(train_state.best_step), verbose=1)
    safe_test(model, data, X_test, y_test)

    tests = [
        (lambda x: x, "x.dat"),
        (lambda x: np.sin(np.pi * x), "sinx.dat"),
        (lambda x: np.sin(2 * np.pi * x), "sin2x.dat"),
        (lambda x: x * np.sin(2 * np.pi * x), "xsin2x.dat"),
    ]
    for u, fname in tests:
        if problem == "ode":
            test_u_ode(nn, system, T, m, model, data, u, fname)

def main():
    # Problems:
    # - "lt": Legendre transform
    # - "ode": Antiderivative, Nonlinear ODE, Gravity pendulum
    # - "dr": Diffusion-reaction
    # - "cvc": Advection
    # - "advd": Advection-diffusion
    problem = "ode"
    T = 1
    system = ode_system(T)

    # Function space
    # space = FinitePowerSeries(N=100, M=1)
    # space = FiniteChebyshev(N=20, M=1)
    # space = GRF(2, length_scale=0.2, N=2000, interp="cubic")  # "lt"
    space = GRF(1, length_scale=0.2, N=1000, interp="cubic")
    # space = GRF(T, length_scale=0.2, N=1000 * T, interp="cubic")

    # Hyperparameters
    m = 100
    num_train = 10000
    num_test = 100000
    lr = 0.001
    epochs = 50000

    # Network
    nn = "opnn"
    activation = "relu"
    initializer = "Glorot normal"  # "He normal" or "Glorot normal"
    dim_x = 1 if problem in ["ode", "lt"] else 2
    if nn == "opnn":
        net = dde.maps.DeepONetCartesianProd(
            [m, 40, 40],
            [dim_x, 40, 40],
            activation,
            initializer,
            # use_bias=True,
            # stacked=False,
        )
    elif nn == "fnn":
        net = dde.maps.FNN([m + dim_x] + [100] * 2 + [1], activation, initializer)
    elif nn == "resnet":
        net = dde.maps.ResNet(m + dim_x, 1, 128, 2, activation, initializer)

    run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test)

if __name__ == "__main__":
    main()

Taking a look briefly at the exports in deepxde.nn.tensorflow I think that the Installation Guide should point to DeepONetCartesianProd and not DeepONet, right? Or am I missing anything? Now even if I make that change, the code is not runnable as is. Do you have any inputs?

I will keep exploring the library, but it'd be great if you could clarify a bit on how to make the above snippet runnable.

Thanks!

bhaveshshrimali commented 3 years ago

Just to add that the above traceback is the same irrespective of whether I run on my local machine or on colab. On my local machine I am running it inside a container built on top of the tensorflow docker image -- with the master branch of deeponet cloned

Thanks in advance. Let me know if you need anymore info!

bhaveshshrimali commented 3 years ago

So if I set the default float to float64, i.e.

dde.config.set_default_float("float64")

I now get a slightly different traceback (still the same error albeit at a different place)

Traceback (most recent call last):
  File "testODE.py", line 143, in <module>
    main()
  File "testODE.py", line 140, in main
    run(problem, system, space, T, m, nn, net, lr, epochs, num_train, num_test)
  File "testODE.py", line 76, in run
    losshistory, train_state = model.train(epochs=epochs, callbacks=[checker])
  File "/usr/local/lib/python3.6/dist-packages/deepxde/utils/internal.py", line 26, in wrapper
    result = f(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/deepxde/model.py", line 353, in train
    self._test()
  File "/usr/local/lib/python3.6/dist-packages/deepxde/model.py", line 489, in _test
    self.train_state.train_aux_vars,
  File "/usr/local/lib/python3.6/dist-packages/deepxde/model.py", line 284, in _outputs_losses
    outs = self.outputs_losses(training, inputs, targets, auxiliary_vars)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 933, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 760, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 3066, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 3463, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 3308, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/func_graph.py", line 1007, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 668, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/func_graph.py", line 994, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:

    /usr/local/lib/python3.6/dist-packages/deepxde/model.py:167 outputs_losses  *
        outputs_ = self.net(inputs, training=training)
    /usr/local/lib/python3.6/dist-packages/deepxde/nn/tensorflow/deeponet.py:67 call  *
        x += self.b
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py:1383 binary_op_wrapper
        raise e
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py:1367 binary_op_wrapper
        return func(x, y, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py:206 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py:1700 _add_dispatch
        return gen_math_ops.add_v2(x, y, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_math_ops.py:465 add_v2
        "AddV2", x=x, y=y, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py:558 _apply_op_helper
        inferred_from[input_arg.type_attr]))

    TypeError: Input 'y' of 'AddV2' Op has type float32 that does not match type float64 of argument 'x'.
bhaveshshrimali commented 3 years ago

So I tried explicitly declaring the type here https://github.com/lululxvi/deepxde/blob/8d417288ad97de074b9645ab29221c0eea4743b9/deepxde/nn/tensorflow/deeponet.py#L48

namely, changing the line to

self.b = tf.Variable(tf.zeros(1, dtype=tf.dtypes.float32))

Doing everything on float64 results in OOM error. (I'm not sure why though, my system does seem to have sufficient memory). Nonetheless, let's see how it progresses.

lululxvi commented 3 years ago

Could you try version 0.11.2?

bhaveshshrimali commented 3 years ago

Could you try version 0.11.2?

Thanks, will try it and report here. Note that I trimmed the problem size and after all the above twitches (a couple more of those tf.cast(...) operations), things do seem to run at least. Will report things with 0.11.2 nontheless.

Thanks!