已合并
add test for add_module and apply api in torch.jit.ScriptModule #38690
bobebest创建于 6月16日
add test for add_module and apply api in torch.jit.ScriptModule #38690
已合并
共 1 个文件变更+115-1
| @@ -2,7 +2,8 @@ | |||
| 2 | Add validation cases for torch.jit.ScriptModule APIs on NPU: | 2 | Add validation cases for torch.jit.ScriptModule APIs on NPU: |
| 3 | 3 | ||
| 4 | PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. | 4 | PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. |
| 5 | -This file validates torch.jit.ScriptModule.bfloat16, torch.jit.ScriptModule.buffers, torch.jit.ScriptModule.children, torch.jit.ScriptModule.code | 5 | +This file validates torch.jit.ScriptModule.bfloat16, torch.jit.ScriptModule.buffers, torch.jit.ScriptModule.children, torch.jit.ScriptModule.code, |
| 6 | +torch.jit.ScriptModule.add_module, torch.jit.ScriptModule.apply | ||
| 6 | """ | 7 | """ |
| 7 | 8 | ||
| 8 | import torch | 9 | import torch |
| @@ -11,6 +12,7 @@ from torch.testing._internal.common_utils import TestCase, run_tests | |||
| 11 | import torch.nn as nn | 12 | import torch.nn as nn |
| 12 | import torch.nn.functional as F | 13 | import torch.nn.functional as F |
| 13 | 14 | ||
| 15 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 14 | 16 | ||
| 15 | class Model(nn.Module): | 17 | class Model(nn.Module): |
| 16 | def __init__(self) -> None: | 18 | def __init__(self) -> None: |
| @@ -104,5 +106,117 @@ class TestJitScriptModuleCode(TestCase): | |||
| 104 | self.assertIn('forward', str(model.code)) | 106 | self.assertIn('forward', str(model.code)) |
| 105 | 107 | ||
| 106 | 108 | ||
| 109 | +class TestScriptModuleAddModule(TestCase): | ||
| 110 | + def test_add_module_on_scriptmodule(self): | ||
| 111 | + class SimpleModel(nn.Module): | ||
| 112 | + def __init__(self): | ||
| 113 | + super().__init__() | ||
| 114 | + self.linear = nn.Linear(5, 3) | ||
| 115 | + def forward(self, x): | ||
| 116 | + return self.linear(x) | ||
| 117 | + | ||
| 118 | + model = SimpleModel() | ||
| 119 | + extra_layer = nn.Linear(3, 2) | ||
| 120 | + model.add_module('sub_module', extra_layer) | ||
| 121 | + | ||
| 122 | + scripted_model = torch.jit.script(model) | ||
| 123 | + self.assertIsInstance(scripted_model, torch.jit.ScriptModule) | ||
| 124 | + | ||
| 125 | + scripted_model.to(device_type) | ||
| 126 | + | ||
| 127 | + self.assertTrue(hasattr(scripted_model, 'sub_module')) | ||
| 128 | + self.assertIsInstance(scripted_model.sub_module, torch.jit.ScriptModule) | ||
| 129 | + self.assertIn('sub_module', scripted_model._modules) | ||
| 130 | + | ||
| 131 | + params = list(scripted_model.sub_module.parameters()) | ||
| 132 | + self.assertEqual(len(params), 2) | ||
| 133 | + self.assertEqual(params[0].shape, (2, 3)) | ||
| 134 | + self.assertEqual(params[0].device.type, 'npu') | ||
| 135 | + self.assertEqual(params[1].device.type, 'npu') | ||
| 136 | + | ||
| 137 | + dummy_input = torch.randn(1, 5, device=device_type) | ||
| 138 | + output = scripted_model(dummy_input) | ||
| 139 | + self.assertEqual(output.shape, (1, 3)) | ||
| 140 | + self.assertEqual(output.device.type, 'npu') | ||
| 141 | + | ||
| 142 | + | ||
| 143 | +class TestScriptModuleApply(TestCase): | ||
| 144 | + def test_apply_modify_parameters(self): | ||
| 145 | + class SimpleModel(nn.Module): | ||
| 146 | + def __init__(self): | ||
| 147 | + super().__init__() | ||
| 148 | + self.conv = nn.Conv2d(1, 2, 3) | ||
| 149 | + self.fc = nn.Linear(2 * 26 * 26, 5) | ||
| 150 | + def forward(self, x): | ||
| 151 | + x = self.conv(x) | ||
| 152 | + x = x.view(x.size(0), -1) | ||
| 153 | + return self.fc(x) | ||
| 154 | + | ||
| 155 | + model = SimpleModel().to(device_type) | ||
| 156 | + scripted = torch.jit.script(model) | ||
| 157 | + self.assertIsInstance(scripted, torch.jit.ScriptModule) | ||
| 158 | + | ||
| 159 | + def zero_params(module): | ||
| 160 | + if hasattr(module, 'weight') and module.weight is not None: | ||
| 161 | + module.weight.data.fill_(0.0) | ||
| 162 | + if hasattr(module, 'bias') and module.bias is not None: | ||
| 163 | + module.bias.data.fill_(0.0) | ||
| 164 | + | ||
| 165 | + scripted.apply(zero_params) | ||
| 166 | + | ||
| 167 | + for param in scripted.parameters(): | ||
| 168 | + self.assertTrue(torch.all(param == 0)) | ||
| 169 | + self.assertEqual(param.device.type, "npu") | ||
| 170 | + | ||
| 171 | + dummy = torch.randn(1, 1, 28, 28, device=device_type) | ||
| 172 | + out = scripted(dummy) | ||
| 173 | + self.assertTrue(torch.all(out == 0)) | ||
| 174 | + self.assertEqual(out.shape, (1, 5)) | ||
| 175 | + self.assertEqual(out.device.type, "npu") | ||
| 176 | + | ||
| 177 | + def test_apply_recursively_visits_all_modules(self): | ||
| 178 | + class Leaf(nn.Module): | ||
| 179 | + def __init__(self): | ||
| 180 | + super().__init__() | ||
| 181 | + self.param = nn.Parameter(torch.randn(2,2)) | ||
| 182 | + def forward(self, x): | ||
| 183 | + return x | ||
| 184 | + | ||
| 185 | + class Container(nn.Module): | ||
| 186 | + def __init__(self): | ||
| 187 | + super().__init__() | ||
| 188 | + self.leaf1 = Leaf() | ||
| 189 | + self.leaf2 = Leaf() | ||
| 190 | + def forward(self, x): | ||
| 191 | + return x | ||
| 192 | + | ||
| 193 | + model = Container().to(device_type) | ||
| 194 | + scripted = torch.jit.script(model) | ||
| 195 | + | ||
| 196 | + visited = set() | ||
| 197 | + def record(module): | ||
| 198 | + visited.add(id(module)) | ||
| 199 | + | ||
| 200 | + scripted.apply(record) | ||
| 201 | + | ||
| 202 | + expected = {id(scripted), id(scripted.leaf1), id(scripted.leaf2)} | ||
| 203 | + self.assertEqual(visited, expected) | ||
| 204 | + | ||
| 205 | + | ||
| 206 | + def test_apply_returns_self(self): | ||
| 207 | + class Dummy(nn.Module): | ||
| 208 | + def forward(self, x): | ||
| 209 | + return x | ||
| 210 | + | ||
| 211 | + model = Dummy().to(device_type) | ||
| 212 | + scripted = torch.jit.script(model) | ||
| 213 | + | ||
| 214 | + def noop(module): | ||
| 215 | + pass | ||
| 216 | + | ||
| 217 | + ret = scripted.apply(noop) | ||
| 218 | + self.assertIs(ret, scripted) | ||
| 219 | + | ||
| 220 | + | ||
| 107 | if __name__ == "__main__": | 221 | if __name__ == "__main__": |
| 108 | run_tests() | 222 | run_tests() |