feat: add eval() in InfiniCoreModule

Author: spike-zhu

python/infinicore/nn/modules/module.py

@@ -19,6 +19,7 @@
 
 _EXTRA_STATE_KEY_SUFFIX = '_extra_state'
 
+T = TypeVar('T', bound='InfiniCoreModule')
 
 class _IncompatibleKeys(namedtuple('IncompatibleKeys', ['missing_keys', 'unexpected_keys'])):
     def __repr__(self):
@@ -514,6 +515,77 @@ def load(module, local_state_dict, prefix=''):
                                self.__class__.__name__, "\n\t".join(error_msgs)))
         return _IncompatibleKeys(missing_keys, unexpected_keys)
 
+    def children(self) -> Iterator['InfiniCoreModule']:
+        r"""Returns an iterator over immediate children modules.
+
+        Yields:
+            Module: a child module
+        """
+        for name, module in self.named_children():
+            yield module
+
+    def named_children(self) -> Iterator[Tuple[str, 'InfiniCoreModule']]:
+        r"""Returns an iterator over immediate children modules, yielding both
+        the name of the module as well as the module itself.
+
+        Yields:
+            (str, Module): Tuple containing a name and child module
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for name, module in model.named_children():
+            >>>     if name in ['conv4', 'conv5']:
+            >>>         print(module)
+
+        """
+        memo = set()
+        for name, module in self._modules.items():
+            if module is not None and module not in memo:
+                memo.add(module)
+                yield name, module
+
+
+    def train(self: T, mode: bool = True) -> T:
+        r"""Sets the module in training mode.
+
+        This has any effect only on certain modules. See documentations of
+        particular modules for details of their behaviors in training/evaluation
+        mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
+        etc.
+
+        Args:
+            mode (bool): whether to set training mode (``True``) or evaluation
+                         mode (``False``). Default: ``True``.
+
+        Returns:
+            Module: self
+        """
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        self.training = mode
+        for module in self.children():
+            module.train(mode)
+        return self
+
+    def eval(self: T) -> T:
+        r"""Sets the module in evaluation mode.
+
+        This has any effect only on certain modules. See documentations of
+        particular modules for details of their behaviors in training/evaluation
+        mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
+        etc.
+
+        This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.
+
+        See :ref:`locally-disable-grad-doc` for a comparison between
+        `.eval()` and several similar mechanisms that may be confused with it.
+
+        Returns:
+            Module: self
+        """
+        return self.train(False)
+
 
     def to(self, device: torch.device) -> "InfiniCoreModule":
         for name, param in self._parameters.items():

test/infinicore/infinicore_nn_test.py

@@ -7,65 +7,90 @@
 # 1. 使用 PyTorch 定义并保存模型
 # ============================================================
 
-class TorchMLP(nn.Module):
-    def __init__(self, in_dim=4, hidden_dim=8, out_dim=2):
+class TorchConvNet(nn.Module):
+    def __init__(self, in_ch=3, hidden_ch=8, out_ch=3):
         super().__init__()
-        self.fc1 = nn.Linear(in_dim, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, out_dim)
+        # 主体网络
+        self.conv1 = nn.Conv2d(in_ch, hidden_ch, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(hidden_ch)
+        self.conv2 = nn.Conv2d(hidden_ch, hidden_ch, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(hidden_ch)
+        self.conv3 = nn.Conv2d(hidden_ch, out_ch, kernel_size=1)
+        self.relu = nn.ReLU()
+
+        # 自定义 Parameter（例如一个可学习缩放因子）
+        self.scale = nn.Parameter(torch.ones(1) * 0.5)
+
+        # 注册一个 buffer（非参数，例如推理时的固定偏移）
+        self.register_buffer("offset", torch.tensor(0.1))
 
     def forward(self, x):
-        x = torch.relu(self.fc1(x))
-        x = self.fc2(x)
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.relu(self.bn2(self.conv2(x)))
+        x = self.conv3(x)
+        # 应用自定义参数和 buffer
+        x = x * self.scale + self.offset
         return x
 
-torch_model = TorchMLP()
-torch_state_dict = torch_model.state_dict()
 
-safetensors.torch.save_file(torch_state_dict, "torch_model.safetensors")
+# ===== 保存 Torch 模型 =====
+torch_model = TorchConvNet()
+torch_state_dict = torch_model.state_dict()
+safetensors.torch.save_file(torch_state_dict, "torch_convnet_with_param.safetensors")
 
 # ============================================================
 # 2. 使用 torch 方式加载并推理
 # ============================================================
 
-torch_model_infer = TorchMLP()
-torch_model_infer.load_state_dict(safetensors.torch.load_file("torch_model.safetensors"))
+torch_model_infer = TorchConvNet()
+torch_model_infer.load_state_dict(safetensors.torch.load_file("torch_convnet_with_param.safetensors"))
 torch_model_infer.eval()
 
-input = torch.rand(1, 4)
+input = torch.rand(1, 3, 8, 8)
 torch_model_out = torch_model_infer(input)
-print("Torch 输出：", torch_model_out.detach().numpy())
+print("Torch 输出：", torch_model_out.detach().numpy().mean())
 
 # ============================================================
 # 3. 使用 InfiniCore.Module 系统加载并推理
 # ============================================================
 
-# ===== 下面定义一个与 TorchMLP 对应的 InfiniCoreModule类 =====
 from python.infinicore.nn.modules import InfiniCoreModule
 
-class InfiniCoreMLP(InfiniCoreModule):
-    def __init__(self, in_dim=4, hidden_dim=8, out_dim=2):
+class InfiniCoreConvNet(InfiniCoreModule):
+    def __init__(self, in_ch=3, hidden_ch=8, out_ch=3):
         super().__init__()
-        self.fc1 = nn.Linear(in_dim, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, out_dim)
+        self.conv1 = nn.Conv2d(in_ch, hidden_ch, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(hidden_ch)
+        self.conv2 = nn.Conv2d(hidden_ch, hidden_ch, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(hidden_ch)
+        self.conv3 = nn.Conv2d(hidden_ch, out_ch, kernel_size=1)
+        self.relu = nn.ReLU()
+
+        # 保持与 Torch 模型一致的自定义参数和 buffer
+        self.scale = nn.Parameter(torch.ones(1) * 0.5)
+        self.register_buffer("offset", torch.tensor(0.1))
 
     def forward(self, x):
-        x = torch.relu(self.fc1(x))
-        x = self.fc2(x)
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.relu(self.bn2(self.conv2(x)))
+        x = self.conv3(x)
+        x = x * self.scale + self.offset
         return x
 
-# ===== 使用 InfiniCoreMLP 读取 safetensors 并推理 =====
-infinicore_model_infer = InfiniCoreMLP()
-infinicore_model_infer.load_state_dict(safetensors.torch.load_file("torch_model.safetensors"))
-infinicore_model_out = infinicore_model_infer.forward(input)
+# ===== 使用 InfiniCoreConvNet 读取 safetensors 并推理 =====
+infinicore_model_infer = InfiniCoreConvNet()
+infinicore_model_infer.load_state_dict(safetensors.torch.load_file("torch_convnet_with_param.safetensors"))
+infinicore_model_infer.eval()
 
-print("InfiniCore 输出：", infinicore_model_out.detach().numpy())
+infinicore_model_out = infinicore_model_infer.forward(input)
+print("InfiniCore 输出：", infinicore_model_out.detach().numpy().mean())
 
 # ============================================================
 # 4. 对比结果
 # ============================================================
 
 diff = (infinicore_model_out - torch_model_out).abs().max().item()
-print(f"InfiniCoreModule 与 Torch 最大误差: {diff:.6f}")
+print(f"InfiniCoreModule 与 Torch 最大误差: {diff:.8f}")
 if diff < 1e-6:
     print("✅ InfiniCoreModule 与 Torch 精度一致!")
 else: