From 7918de172ba682794f028dd33f5d915c049b0f0b Mon Sep 17 00:00:00 2001
From: xiaoyewww <641311428@qq.com>
Date: Thu, 26 Dec 2024 00:30:38 +0800
Subject: [PATCH] fix(neuralop): init uno model weights as torch

---
 neuralop/models/uno.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

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)