issue/890 - 为python端的nn.module添加to函数

Author: pengcheng888

python/infinicore/nn/modules/module.py

@@ -32,6 +32,7 @@
 
 import infinicore
 
+from ...device import device as InfiniCoreDevice
 from ...tensor import Tensor
 from ..parameter import InfiniCoreParameter as Parameter
 
@@ -481,15 +482,14 @@ def _load_from_state_dict(
                         f"While copying the parameter named {key}, expected Tensor from checkpoint but received {type(input_param)}"
                     )
 
-                if (
-                    (param.shape == input_param.shape)
-                    and (param.dtype == input_param.dtype)
-                    and (param.device == input_param.device)
+                if (param.shape == input_param.shape) and (
+                    param.dtype == input_param.dtype
                 ):
                     param.copy_(input_param)
                 else:
-                    print(f"param '{name}' don't match input_param '{key}'")
-                    setattr(self, name, input_param)
+                    raise KeyError(
+                        f"param '{name}' don't match input_param '{key}' with shape or dtype"
+                    )
 
             elif strict:
                 missing_keys.append(key)
@@ -848,10 +848,29 @@ def eval(self: T) -> T:
         Returns:
             Module: self
         """
-        pass
+        raise KeyError("not support")
 
     def _apply(self, fn, recurse=True):
-        raise KeyError("not support")
+        if recurse:
+            for module in self.children():
+                module._apply(fn)
 
-    def to(self, *args, **kwargs):
-        raise KeyError("not support")
+        for key, param in self._parameters.items():
+            if param is not None:
+                setattr(self, key, fn(param))
+
+        for key, buf in self._buffers.items():
+            if buf is not None:
+                setattr(self, key, fn(buf))
+
+        return self
+
+    def to(self, device: str | InfiniCoreDevice):
+        if device is None:
+            raise ValueError("device cannot be None")
+        device = InfiniCoreDevice(device)
+
+        def convert(t):
+            return t.to(device)
+
+        return self._apply(convert)

test/infinicore/nn/module.py

@@ -44,6 +44,7 @@ def __init__(self):
     def forward(self):
         return infinicore.add(self.a, self.b)
 
+
 infinicore_model_infer = InfiniCoreNet()
 # ============================================================
 # 2. 加载权重
@@ -75,6 +76,98 @@ def forward(self):
 
 
 # ============================================================
-# 5. to测试，buffer测试
+# 5. to测试 - 测试模型在不同设备间的转换
 # ============================================================
-# 等待添加
+print("\n" + "=" * 60)
+print("5. to测试 - 设备转换测试")
+print("=" * 60)
+
+# 5.1 打印初始状态
+print("\n5.1 初始状态:")
+print("-" * 40)
+print("Parameters:")
+for name, param in infinicore_model_infer.named_parameters():
+    print(f"  {name}: shape={param.shape}, dtype={param.dtype}, device={param.device}")
+print("Buffers:")
+buffers_exist = False
+for name, buf in infinicore_model_infer.named_buffers():
+    buffers_exist = True
+    print(f"  {name}: shape={buf.shape}, dtype={buf.dtype}, device={buf.device}")
+if not buffers_exist:
+    print("  (无buffers)")
+
+# 5.2 测试转换到CUDA设备（使用device对象）
+print("\n5.2 转换到CUDA设备 (使用 infinicore.device('cuda', 0)):")
+print("-" * 40)
+target_device_cuda = infinicore.device("cuda", 0)
+infinicore_model_infer.to(target_device_cuda)
+
+print("转换后的Parameters:")
+for name, param in infinicore_model_infer.named_parameters():
+    print(f"  {name}: shape={param.shape}, dtype={param.dtype}, device={param.device}")
+    # 验证设备是否正确转换
+    assert param.device == target_device_cuda, (
+        f"参数 {name} 的设备转换失败: 期望 {target_device_cuda}, 实际 {param.device}"
+    )
+if buffers_exist:
+    print("转换后的Buffers:")
+    for name, buf in infinicore_model_infer.named_buffers():
+        print(f"  {name}: shape={buf.shape}, dtype={buf.dtype}, device={buf.device}")
+        assert buf.device == target_device_cuda, (
+            f"Buffer {name} 的设备转换失败: 期望 {target_device_cuda}, 实际 {buf.device}"
+        )
+print("✓ CUDA设备转换验证通过")
+
+# 5.3 测试转换到CPU设备（使用device对象）
+print("\n5.3 转换到CPU设备 (使用 infinicore.device('cpu', 0)):")
+print("-" * 40)
+target_device_cpu = infinicore.device("cpu", 0)
+infinicore_model_infer.to(target_device_cpu)
+
+print("转换后的Parameters:")
+for name, param in infinicore_model_infer.named_parameters():
+    print(f"  {name}: shape={param.shape}, dtype={param.dtype}, device={param.device}")
+    # 验证设备是否正确转换
+    assert param.device == target_device_cpu, (
+        f"参数 {name} 的设备转换失败: 期望 {target_device_cpu}, 实际 {param.device}"
+    )
+if buffers_exist:
+    print("转换后的Buffers:")
+    for name, buf in infinicore_model_infer.named_buffers():
+        print(f"  {name}: shape={buf.shape}, dtype={buf.dtype}, device={buf.device}")
+        assert buf.device == target_device_cpu, (
+            f"Buffer {name} 的设备转换失败: 期望 {target_device_cpu}, 实际 {buf.device}"
+        )
+print("✓ CPU设备转换验证通过")
+
+# 5.4 测试使用字符串参数转换到CUDA设备
+print("\n5.4 转换到CUDA设备 (使用字符串 'cuda'):")
+print("-" * 40)
+infinicore_model_infer.to("cuda")
+
+print("转换后的Parameters:")
+for name, param in infinicore_model_infer.named_parameters():
+    print(f"  {name}: shape={param.shape}, dtype={param.dtype}, device={param.device}")
+    # 验证设备是否正确转换（字符串'cuda'会被转换为cuda设备）
+    assert param.device.type == "cuda", (
+        f"参数 {name} 的设备转换失败: 期望 cuda, 实际 {param.device.type}"
+    )
+if buffers_exist:
+    print("转换后的Buffers:")
+    for name, buf in infinicore_model_infer.named_buffers():
+        print(f"  {name}: shape={buf.shape}, dtype={buf.dtype}, device={buf.device}")
+        assert buf.device.type == "cuda", (
+            f"Buffer {name} 的设备转换失败: 期望 cuda, 实际 {buf.device.type}"
+        )
+print("✓ 字符串参数设备转换验证通过")
+
+# 5.5 验证to方法返回self（链式调用支持）
+print("\n5.5 测试to方法的返回值（链式调用）:")
+print("-" * 40)
+result = infinicore_model_infer.to(infinicore.device("cpu", 0))
+assert result is infinicore_model_infer, "to方法应该返回self以支持链式调用"
+print("✓ to方法返回值验证通过")
+
+print("\n" + "=" * 60)
+print("所有to测试通过！")
+print("=" * 60 + "\n")