diff --git a/neuralop/models/uno.py b/neuralop/models/uno.py index f089295..f2a0065 100644 --- a/neuralop/models/uno.py +++ b/neuralop/models/uno.py @@ -1,6 +1,8 @@ +import math import paddle import paddle.nn as nn import paddle.nn.functional as F +import ppsci from ..layers.fno_block import FNOBlocks from ..layers.mlp import MLP @@ -10,6 +12,19 @@ from ..layers.spectral_convolution import SpectralConv +def kaiming_init(layer): + if isinstance(layer, (nn.layer.conv._ConvNd, nn.Linear)): + print(f"layer: {layer} ") + init_kaimingUniform = paddle.nn.initializer.KaimingUniform(nonlinearity='leaky_relu', negative_slope=math.sqrt(5)) + init_kaimingUniform(layer.weight) + if layer.bias is not None: + fan_in, _ = ppsci.utils.initializer._calculate_fan_in_and_fan_out(layer.weight) + if fan_in != 0: + bound = 1 / math.sqrt(fan_in) + init_uniform = paddle.nn.initializer.Uniform(low=-bound, high=bound) + init_uniform(layer.bias) + + class UNO(nn.Layer): """U-Shaped Neural Operator [1]_ @@ -279,6 +294,8 @@ def __init__( non_linearity=non_linearity, ) + self.apply(kaiming_init) + def forward(self, x, **kwargs): x = self.lifting(x)