已合并
test(nn): add test for models api:torch.nn.Module, torch.nn.Module.state_dict, torch.nn.ModuleDict, torch.nn.ModuleList #31787
dinglaiping创建于 3月13日
test(nn): add test for models api:torch.nn.Module, torch.nn.Module.state_dict, torch.nn.ModuleDict, torch.nn.ModuleList #31787
已合并
从已删除 :addtest-models-api-2.7.1合入到Ascend/pytorchv2.7.1
共 1 个文件变更+83-0
| @@ -0,0 +1,83 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.nn.models APIs on NPU: | ||
| 3 | +1. test/test_modules.py from PyTorch community lacks sufficient API validations, so this file is added. | ||
| 4 | +2. This file validates torch.nn.Module, torch.nn.Module.state_dict, torch.nn.ModuleDict, torch.nn.ModuleList (extendable). | ||
| 5 | +""" | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +import torch.nn as nn | ||
| 9 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 10 | +import torch_npu | ||
| 11 | + | ||
| 12 | +device = torch.device("npu:0") | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +class TestNPUModuleAPIs(TestCase): | ||
| 16 | + def test_module_device_consistency(self): | ||
| 17 | + """验证Module设备迁移后参数/缓冲区设备一致""" | ||
| 18 | + m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20)).to(device) | ||
| 19 | + | ||
| 20 | + for p in m.parameters(): | ||
| 21 | + self.assertEqual(p.device, device) | ||
| 22 | + for b in m.buffers(): | ||
| 23 | + self.assertEqual(b.device, device) | ||
| 24 | + | ||
| 25 | + def test_module_state_dict(self): | ||
| 26 | + """验证state_dict保存/加载后状态一致""" | ||
| 27 | + base = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20)) | ||
| 28 | + m = nn.Module() | ||
| 29 | + m.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base}) | ||
| 30 | + m.list = nn.ModuleList([nn.Linear(20, 30), base]) | ||
| 31 | + m.to(device) | ||
| 32 | + | ||
| 33 | + sd = m.state_dict() | ||
| 34 | + torch.npu.synchronize() | ||
| 35 | + | ||
| 36 | + for v in sd.values(): | ||
| 37 | + self.assertEqual(v.device, device) | ||
| 38 | + | ||
| 39 | + base2 = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.BatchNorm1d(20)) | ||
| 40 | + m2 = nn.Module() | ||
| 41 | + m2.dict = nn.ModuleDict({"linear": nn.Linear(10, 20), "base": base2}) | ||
| 42 | + m2.list = nn.ModuleList([nn.Linear(20, 30), base2]) | ||
| 43 | + m2.load_state_dict(sd) | ||
| 44 | + m2.to(device) | ||
| 45 | + torch.npu.synchronize() | ||
| 46 | + | ||
| 47 | + for (n1, p1), (n2, p2) in zip(m.named_parameters(), m2.named_parameters()): | ||
| 48 | + self.assertTrue(torch.allclose(p1, p2)) | ||
| 49 | + | ||
| 50 | + def test_moduledict_operations(self): | ||
| 51 | + """验证ModuleDict增删/索引/遍历""" | ||
| 52 | + m = nn.ModuleDict({"a": nn.Linear(10, 20).to(device)}) | ||
| 53 | + | ||
| 54 | + self.assertIn("a", m) | ||
| 55 | + m["b"] = nn.Linear(20, 30).to(device) | ||
| 56 | + self.assertEqual(m["b"].weight.device, device) | ||
| 57 | + del m["b"] | ||
| 58 | + self.assertNotIn("b", m) | ||
| 59 | + | ||
| 60 | + for sub in m.values(): | ||
| 61 | + for p in sub.parameters(): | ||
| 62 | + self.assertEqual(p.device, device) | ||
| 63 | + | ||
| 64 | + def test_modulelist_operations(self): | ||
| 65 | + """验证ModuleList索引/新增/删除/遍历""" | ||
| 66 | + m = nn.ModuleList([nn.Linear(20, 30).to(device), nn.BatchNorm1d(30).to(device)]) | ||
| 67 | + | ||
| 68 | + self.assertEqual(m[0].weight.device, device) | ||
| 69 | + self.assertEqual(m[1].running_mean.device, device) | ||
| 70 | + | ||
| 71 | + m.append(nn.Linear(30, 40).to(device)) | ||
| 72 | + m.insert(0, nn.Linear(10, 20).to(device)) | ||
| 73 | + m.pop(-1) # 修复：指定索引 | ||
| 74 | + m.pop(0) # 修复：指定索引 | ||
| 75 | + | ||
| 76 | + self.assertEqual(len(m), 2) | ||
| 77 | + for sub in m: | ||
| 78 | + for p in sub.parameters(): | ||
| 79 | + self.assertEqual(p.device, device) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +if __name__ == "__main__": | ||
| 83 | + run_tests() | ||