What's wrong with my implementation of MsFFN

[ghost]

Hi, I tried to implement MsFFN using pytorch, but the result is not as expected. The following is my code.

import deepxde as dde
import numpy as np

from deepxde.backend import torch
from deepxde.nn import activations, initializers, NN
from deepxde import config


class MsFFN(NN):
    def __init__(
        self,
        layer_sizes,
        activation,
        kernel_initializer,
        sigmas,
        regularization=None,
        use_bias=True,
    ):
        super().__init__()
        self.activation = activations.get(activation)
        initializer = initializers.get(kernel_initializer)
        initializer_zero = initializers.get("zeros")

        self.layer_size = layer_sizes
        self.sigmas = sigmas  # list or tuple
        self.fourier_feature_weights = None

        self.linears = torch.nn.ModuleList()
        for i in range(1, len(layer_sizes)-2):
            self.linears.append(
                torch.nn.Linear(
                    layer_sizes[i], layer_sizes[i+1], dtype=config.real(torch)
                )
            )
            initializer(self.linears[-1].weight)
            initializer_zero(self.linears[-1].bias)
        self.linears.append(torch.nn.Linear(layer_sizes[-2]*len(self.sigmas), layer_sizes[-1]))
        initializer(self.linears[-1].weight)
        initializer_zero(self.linears[-1].bias)


    def _fourier_feature_forward(self, y, sigma):
        b = torch.randn(y.size(1), self.layer_size[1] // 2) * sigma
        y = torch.cat(
            [
                torch.cos(torch.matmul(y, b)),
                torch.sin(torch.matmul(y, b)),
            ],
            dim=1,
        )
        return y, b

    def _fully_connected_forward(self, y):
        for linear in self.linears[:-1]:
            y = self.activation(linear(y))
        return y
        
    def forward(self, inputs):
        y = inputs
        # batch_size = y.size(0)
        if self._input_transform is not None:
            y = self._input_transform(y)
        
        yb = [self._fourier_feature_forward(y, sigma) for sigma in self.sigmas]
        y = [elem[0] for elem in yb]
        self.fourier_feature_weights = [elem[1] for elem in yb]
        # fully-connected layers (reuse)
        y = [self._fully_connected_forward(_y) for _y in y]

        # concatenate all the fourier features
        y = torch.cat(y, dim=1)
        y = self.linears[-1](y)

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

I run an example as follows.

def pde(x, y):
    dy_xx = dde.grad.hessian(y, x)
    return dy_xx + 4*np.pi**2*torch.sin(2*np.pi*x) + 250*np.pi**2*torch.sin(50*np.pi*x)
    


def boundary(x, on_boundary):
    return on_boundary


def func(x):
    return np.sin(2*np.pi * x) + 0.1*np.sin(50*np.pi*x)


geom = dde.geometry.Interval(0, 1)
bc = dde.icbc.DirichletBC(geom, func, boundary)
data = dde.data.PDE(geom, pde, bc, 500, 2, solution=func, num_test=100)

layer_size = [1, 100, 100] + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = MsFFN(layer_size, activation, initializer, [1, 50])

model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])

losshistory, train_state = model.train(iterations=5000)

dde.saveplot(losshistory, train_state, issave=True, isplot=True)

The result is

0         [5.93e+06, 2.25e-01]    [1.32e+07, 2.98e-01]    [1.27e+00]    
1000      [3.00e+06, 2.13e+01]    [2.91e+06, 2.66e+01]    [7.56e+00]    
2000      [3.03e+06, 3.05e+01]    [3.15e+06, 3.54e+01]    [8.56e+00]    
3000      [3.07e+06, 1.81e+01]    [3.02e+06, 1.59e+01]    [4.88e+00]    
4000      [2.91e+06, 1.36e+01]    [2.99e+06, 1.56e+01]    [4.03e+00]    
5000      [3.05e+06, 1.81e+01]    [3.00e+06, 1.20e+01]    [2.41e+00]    

Best model at step 4000:
  train loss: 2.91e+06
  test loss: 2.99e+06
  test metric: [4.03e+00]

While I get the expected result when I using the tensorflow.compat.v1 backend.

import deepxde as dde
import numpy as np
from deepxde.backend import tf

def pde(x, y):
    dy_xx = dde.grad.hessian(y, x)
    return dy_xx + 4*np.pi**2*tf.sin(2*np.pi*x) + 250*np.pi**2*tf.sin(50*np.pi*x)
    


def boundary(x, on_boundary):
    return on_boundary


def func(x):
    return np.sin(2*np.pi * x) + 0.1*np.sin(50*np.pi*x)


geom = dde.geometry.Interval(0, 1)
bc = dde.icbc.DirichletBC(geom, func, boundary)
data = dde.data.PDE(geom, pde, bc, 500, 2, solution=func, num_test=100)

layer_size = [1, 100, 100] + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.MsFFN(layer_size, activation, initializer, [1, 50])

model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])

losshistory, train_state = model.train(iterations=5000)

dde.saveplot(losshistory, train_state, issave=True, isplot=True)

Step      Train loss              Test loss               Test metric   
0         [6.81e+06, 1.88e-02]    [6.86e+06, 1.88e-02]    [1.18e+00]    
1000      [4.20e+01, 5.63e+00]    [4.87e+01, 5.63e+00]    [1.94e+00]    
2000      [1.17e+01, 1.08e-01]    [1.34e+01, 1.08e-01]    [2.68e-01]    
3000      [4.31e+00, 1.92e-04]    [5.06e+00, 1.92e-04]    [1.04e-02]    
4000      [6.69e+00, 1.65e-05]    [6.98e+00, 1.65e-05]    [2.64e-03]    
5000      [3.81e+01, 3.57e-06]    [4.05e+01, 3.57e-06]    [1.98e-03]    

Best model at step 3000:
  train loss: 4.31e+00
  test loss: 5.06e+00
  test metric: [1.04e-02]

[lululxvi]

Would the following line generates a new random b every time?

def _fourier_feature_forward(self, y, sigma):
    b = torch.randn(y.size(1), self.layer_size[1] // 2) * sigma