Small loss but big test error in backend tf with Darcy flow.

[BraveDrXuTF]

Dear Dr. Lu and contributors,
Hi, I'm handling Darcy flow case with Pi-deeponet in tf1 backend, with $\nabla \cdot k(\mathbf{x}) \nabla T=Q $, where $\mathbf{x}$ is piecewise constant, where $k'(\mathbf{x})=0 $, so it would not be affected by the issue (and no definition on the boundaries of different constants, so I'm also having doubt whether it can be applied with Pi-deeponet directly). And here is PDE definition of my script (If needed I could paste all, but they are messy now):

def pde(x, y, v):  
    # v is the aux tensor, k
    return dde.grad.jacobian(v*dde.grad.jacobian(y, x, i=0, j=0),x,i=0,j=0) + dde.grad.jacobian(v*dde.grad.jacobian(y, x, i=0, j=1),x,i=0,j=1) - 1

where I got the following training log:

Step      Train loss              Test loss               Test metric
0         [9.78e-01, 8.38e-21]    [1.34e+00, 1.77e-20]    []  
1000      [3.74e-01, 2.06e-19]    [3.45e-01, 2.06e-19]    []  
2000      [2.98e-01, 2.37e-19]    [2.86e-01, 2.37e-19]    []  
3000      [2.96e-01, 2.78e-19]    [2.50e-01, 2.78e-19]    []  
4000      [3.02e-01, 2.53e-19]    [2.46e-01, 2.53e-19]    []  

Please ignore the test loss as I have not treated it well. According to the source code in deepxde,
self.outputs_losses_train = [self.net.outputs, losses_train],
As you can see, the first col of training loss is not loss, instead, what displayed is essentially the net output and real train loss. I'm really confused by this. losses_train is so small. Do you think it is the problem of my pde definition?
And the test loss is really high, around 2.-3., even worse than no training.
In fact, I'm working to compare PI-deepoent with my variational operator learning method on the Darcy flow case. Any advice would be helpful.

[lululxvi]

According to the source code in deepxde, self.outputs_losses_train = [self.net.outputs, losses_train], As you can see, the first col of training loss is not loss, instead, what displayed is essentially the net output and real train loss.

It should be the loss. Could you provide a link to the code?

[BraveDrXuTF]

Of course, it is in model.py.

[BraveDrXuTF]

And this is the deeponet tensorflow v1, which has overridden the outputs function from the parent class.

class DeepONetCartesianProd(NN):
    """Deep operator network for dataset in the format of Cartesian product.

    Args:
        layer_size_branch: A list of integers as the width of a fully connected network,
            or `(dim, f)` where `dim` is the input dimension and `f` is a network
            function. The width of the last layer in the branch and trunk net should be
            equal.
        layer_size_trunk (list): A list of integers as the width of a fully connected
            network.
        activation: If `activation` is a ``string``, then the same activation is used in
            both trunk and branch nets. If `activation` is a ``dict``, then the trunk
            net uses the activation `activation["trunk"]`, and the branch net uses
            `activation["branch"]`.
    """

    def __init__(
        self,
        layer_size_branch,
        layer_size_trunk,
        activation,
        kernel_initializer,
        regularization=None,
    ):
        super().__init__()
        self.layer_size_func = layer_size_branch
        self.layer_size_loc = layer_size_trunk
        if isinstance(activation, dict):
            self.activation_branch = activations.get(activation["branch"])
            self.activation_trunk = activations.get(activation["trunk"])
        else:
            self.activation_branch = self.activation_trunk = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.regularizer = regularizers.get(regularization)

        self._inputs = None

    @property
    def inputs(self):
        return self._inputs

    @property
    def outputs(self):
        return self.y

    @property
    def targets(self):
        return self.target

    @timing
    def build(self):
        print("Building DeepONetCartesianProd...")
        self.X_func = tf.placeholder(config.real(tf), [None, self.layer_size_func[0]])
        self.X_loc = tf.placeholder(config.real(tf), [None, self.layer_size_loc[0]])
        self._inputs = [self.X_func, self.X_loc]

        # Branch net to encode the input function
        y_func = self.X_func
        if callable(self.layer_size_func[1]):
            # User-defined network
            y_func = self.layer_size_func[1](y_func)
        else:
            # Fully connected network
            for i in range(1, len(self.layer_size_func) - 1):
                y_func = tf.layers.dense(
                    y_func,
                    self.layer_size_func[i],
                    activation=self.activation_branch,
                    kernel_initializer=self.kernel_initializer,
                    kernel_regularizer=self.regularizer,
                )
            y_func = tf.layers.dense(
                y_func,
                self.layer_size_func[-1],
                kernel_initializer=self.kernel_initializer,
                kernel_regularizer=self.regularizer,
            )

        # Trunk net to encode the domain of the output function
        y_loc = self.X_loc
        if self._input_transform is not None:
            y_loc = self._input_transform(y_loc)
        for i in range(1, len(self.layer_size_loc)):
            y_loc = tf.layers.dense(
                y_loc,
                self.layer_size_loc[i],
                activation=self.activation_trunk,
                kernel_initializer=self.kernel_initializer,
                kernel_regularizer=self.regularizer,
            )

        # Dot product
        if y_func.shape[-1] != y_loc.shape[-1]:
            raise AssertionError(
                "Output sizes of branch net and trunk net do not match."
            )
        self.y = tf.einsum("bi,ni->bn", y_func, y_loc)
        # Add bias
        b = tf.Variable(tf.zeros(1, dtype=config.real(tf)))
        self.y += b

        if self._output_transform is not None:
            self.y = self._output_transform(self._inputs, self.y)

        self.target = tf.placeholder(config.real(tf), [None, None])
        self.built = True

[BraveDrXuTF]

And here is the full code:

import sys

import matplotlib.pyplot as plt
import numpy as np
from scipy import fft, io
from sklearn.preprocessing import StandardScaler
from deepxde.data.function_spaces import FunctionSpace
from deepxde.data.sampler import BatchSampler
import deepxde as dde
from deepxde.backend import tf, load_backend,get_preferred_backend
from scipy.interpolate import interpn

import os

class NonAnalyticFunctionspace(FunctionSpace):
    r"""
    Non analytic function space.

    """

    def __init__(self, len_train_data, resolution, 
                 data_dir=r"E:\xtf\data\darcyflow\neuraloperator-master\data_generation\darcy\data\512",shift=0,shuffle=True, feature_name="input_darcy_512_"):

        # self.maxiter = maxiter
        # self.i = 0
        self.trainSampler = BatchSampler(len_train_data, shuffle=shuffle)

        self.data_dir = data_dir

        self.resolution = resolution

        self.parent_resolution = 512

        self.parent_grid = (np.linspace(0, 1, self.parent_resolution), np.linspace(0, 1, self.parent_resolution))

        self.son_grid = np.meshgrid(np.linspace(0, 1, self.resolution), np.linspace(0, 1, self.resolution), indexing='ij')

        self.shift = shift
        
        self.feature_name = feature_name 

        # self.interp = interp(self.parent_grid, )
    def random(self, bssize):
        # input random feature from file

        indices = self.trainSampler.get_next(bssize)

        # get the file 
        # indice_list = [interpn(self.parent_grid, np.load(os.path.join(self.data_dir, r"{indices[i]}.npy")),self.son_grid) for i in 
        # indices]
        # feature_list = [np.load(os.path.join(self.data_dir, r"input_darcy_512_{}.npy".format(i+1+self.shift))) for i in 
        # indices]
        feature_list = [np.squeeze(np.load(os.path.join(self.data_dir, self.feature_name+"{}.npy".format(i+1+self.shift)))) for i in 
        indices]
        return np.stack(feature_list, axis=2)

    def eval_one(self, feature, x):
        
        pass

    def eval_batch(self, features, xs):
        # interp the feature, get target value
        
        return interpn(self.parent_grid, features, xs).transpose()


def u_dirichlet(_):

    return 0


def boundary(_, on_boundary):
    
    return on_boundary


def pde(x, y, v):

    
    return dde.grad.jacobian(v*dde.grad.jacobian(y, x, i=0, j=0),x,i=0,j=0) + dde.grad.jacobian(v*dde.grad.jacobian(y, x, i=0, j=1),x,i=0,j=1) - 1



def main(config):
    # x_train, y_train = get_data(r"data\darcy_rectangular_pwc\piececonst_r421_N1024_smooth1.mat", 1000)
    # x_test, y_test = get_data("data\darcy_rectangular_pwc\piececonst_r421_N1024_smooth2.mat", 200)
    # data = dde.data.TripleCartesianProd(x_train, y_train, x_test, y_test)
    
    geom = dde.geometry.Rectangle([0, 0], [1, 1])
    
    dirichletbc = dde.icbc.DirichletBC(geom, u_dirichlet, boundary)
     
    pde = dde.data.PDE(geom, config.get('pde'), dirichletbc, num_domain=1200, num_boundary=124,train_distribution='uniform')
    
    nona_space = NonAnalyticFunctionspace(len_train_data=config.get('trainnum'), resolution=config.get('resolution'),shift=2010)

    evaluation_points = np.linspace(0,1, config.get('resolution'))

    evaluation_pointsx, evaluation_pointsxy = np.meshgrid(evaluation_points, evaluation_points,indexing='ij')

    evaluation_points = np.vstack((evaluation_pointsx.ravel(), evaluation_pointsxy.ravel())).T

    m = config.get('resolution') ** 2
    pde_op = dde.data.PDEOperatorCartesianProd(
        pde,
        nona_space,
        evaluation_points,
        num_function=config['trainnum'],
        batch_size=16,
        num_test=1,
    )

    # m = config.get('resolution') ** 2
    activation = "relu"
    branch = tf.keras.Sequential(
        [
            tf.keras.layers.InputLayer(input_shape=(m,)),
            tf.keras.layers.Reshape((config.get('resolution'), config.get('resolution'), 1)),
            tf.keras.layers.Conv2D(4, (2, 2), strides=1, activation=activation),
            tf.keras.layers.Conv2D(16, (3, 3), strides=2, activation=activation),
            tf.keras.layers.Conv2D(16, (3, 3), strides=2, activation=activation),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(256, activation=activation),
            tf.keras.layers.Dense(128),
        ]
    )
    branch.summary()
    activation = "relu"
    net = dde.maps.DeepONetCartesianProd(
        [m, branch], [2, 128, 128, 128, 128], activation, "Glorot normal"
    )

    def output_transform(inputs, outputs):
        # shape become unknown after this transform, I use einsum to solve the problem, making sure tensorflow keeps aware of the shape.
        return 0.01 * tf.einsum("bn,n->bn" , outputs , tf.squeeze(tf.sin(np.pi*inputs[1][:, 0:1])*tf.sin(np.pi*inputs[1][:, 1:2])))
        # return tf.scalar_mul(0.001 , tf.scalar_mul(outputs , tf.squeeze(tf.sin(np.pi*inputs[1][:, 0:1])*tf.sin(np.pi*inputs[1][:, 1:2]))))
        # return x[:, 0:1]**2 * tf.cos(np.pi * x[:, 0:1]) + x[:, 1:2] * (1 - x[:, 0:1]**2) * y
    # # mollifier = torch.sin(np.pi * mesh[..., 0]) * torch.sin(np.pi * mesh[..., 1]) * 0.001
    # # mollifier = mollifier.to(rank)
    net.apply_output_transform(output_transform)

    # Define and train model
    model = dde.Model(pde_op, net)

    # adamwopt=extend_with_decoupled_weight_decay(tf.keras.optimizers.Adam)
    # opt_instance = adamwopt(learning_rate=0.001, weight_decay=1e-4)

    # MyAdamW = extend_with_decoupled_weight_decay(tf.train.AdamOptimizer)
    # optimizer = MyAdamW(weight_decay=0.0001, learning_rate= 0.02)
    # decay_step = config['iteration_num'] - config['iteration_num']//10
    # warm_step = config['iteration_num']//10
    model.compile(
        "adam",
        lr=0.02,
        # decay=("cosinewarm",  decay_step , warm_step ),
        # metrics=["mean l2 relative error"],
    )

    losshistory, train_state = model.train(iterations=config['iteration_num'], batch_size=config['batch_size'])

    dde.utils.plot_loss_history(losshistory)
    plt.show()




    nona_test_xspace = NonAnalyticFunctionspace(len_train_data=config.get('testnum'), resolution=config.get('resolution'),shift=10,shuffle=False)

    nona_test_yspace = NonAnalyticFunctionspace(len_train_data=config.get('testnum'), resolution=config.get('resolution'),shift=10,data_dir=r'E:\xtf\data\darcyflow\neuraloperator-master\data_generation\darcy\data\512_results\transpose',shuffle=False,feature_name="")
    test_err = 0
    for i in range(config.get('testnum')):
        test_feature_x = nona_test_xspace.random(1)
        test_feature_y = nona_test_yspace.random(1)

        test_feature_x = nona_test_xspace.eval_batch(test_feature_x,evaluation_points)
        test_feature_y = nona_test_yspace.eval_batch(test_feature_y,evaluation_points)


        evaluation_points = np.linspace(0,1, config.get('resolution'))

        evaluation_pointsx, evaluation_pointsxy = np.meshgrid(evaluation_points, evaluation_points,indexing='ij')

        evaluation_points = np.vstack((evaluation_pointsx.ravel(), evaluation_pointsxy.ravel())).T


        test_pred = model.predict((test_feature_x, evaluation_points))
        test_pred = test_pred.reshape((1,config.get('resolution'),config.get('resolution')))
        test_feature_y = test_feature_y.reshape((1,config.get('resolution'),config.get('resolution')))
        test_err += dde.metrics.l2_relative_error(test_pred,test_feature_y)
        if i%100==0:
            print('test error: ', dde.metrics.l2_relative_error(test_pred,test_feature_y))
    test_err = test_err/config.get('testnum')
    print('ave test error: ', test_err)
    
    
    
    # test_pred = 




            # func_feats = self.func_space.random(self.num_func)
            # func_vals = self.func_space.eval_batch(func_feats, self.eval_pts)
            # vx = self.func_space.eval_batch(
            #     func_feats, self.pde.train_x[:, self.func_vars]
            # )
if __name__ == "__main__":
    config={
        # 'trainresolution': 64,
        'resolution': 64,
        'trainnum': 2000,
        'testnum': 2000,
        'pde': pde,
        'batch_size': 16,
        'iteration_num': 12500,
        
    }
    main(config)

[BraveDrXuTF]

In this code, I write a function space to downsample my data from npy file. If needed I can provide data, but they are big.

[lululxvi]

No. The output in the screen is only the loss.