已合并
test(jit):add test fot ignore #39838
wei-pengfei22创建于 7月2日
test(jit):add test fot ignore #39838
已合并
共 1 个文件变更+162-0
| @@ -14,8 +14,12 @@ torch.jit.ScriptModule.register_state_dict_post_hook | |||
| 14 | (extendable) | 14 | (extendable) |
| 15 | """ | 15 | """ |
| 16 | 16 | ||
| 17 | +import warnings | ||
| 18 | +import io | ||
| 19 | +import numpy as np | ||
| 17 | import torch | 20 | import torch |
| 18 | import torch.nn as nn | 21 | import torch.nn as nn |
| 22 | +from torch.testing._internal.jit_utils import JitTestCase | ||
| 19 | from torch.testing._internal.common_utils import run_tests, TestCase | 23 | from torch.testing._internal.common_utils import run_tests, TestCase |
| 20 | 24 | ||
| 21 | 25 | ||
| @@ -522,5 +526,163 @@ class TestScriptModuleHooks(TestCase): | |||
| 522 | self.assertEqual(len(called), 1) | 526 | self.assertEqual(len(called), 1) |
| 523 | 527 | ||
| 524 | 528 | ||
| 529 | +class TestJitIgnoreNPU(JitTestCase): | ||
| 530 | + | ||
| 531 | + def getExportImportCopy(self, mod): | ||
| 532 | + buffer = io.BytesIO() | ||
| 533 | + torch.jit.save(mod, buffer) | ||
| 534 | + buffer.seek(0) | ||
| 535 | + return torch.jit.load(buffer) | ||
| 536 | + | ||
| 537 | + def test_ignore_decorator(self): | ||
| 538 | + with warnings.catch_warnings(record=True) as warns: | ||
| 539 | + warnings.simplefilter("always") | ||
| 540 | + | ||
| 541 | + class M(nn.Module): | ||
| 542 | + def __init__(self) -> None: | ||
| 543 | + super().__init__() | ||
| 544 | + self.state = nn.Buffer(torch.zeros(1).to(device_type)) | ||
| 545 | + | ||
| 546 | + def forward(self, x: torch.Tensor) -> torch.Tensor: | ||
| 547 | + return x | ||
| 548 | + | ||
| 549 | + | ||
| 550 | + def ignored_func(self, x: torch.Tensor) -> torch.Tensor: | ||
| 551 | + self.state = torch.tensor([999.0]).to(device_type) | ||
| 552 | + return x * 10 | ||
| 553 | + | ||
| 554 | + raw_m = M().to(device_type) | ||
| 555 | + m = torch.jit.script(raw_m) | ||
| 556 | + | ||
| 557 | + target_warns = [w for w in warns if "TorchScript will now drop the function" in str(w.message)] | ||
| 558 | + self.assertEqual(len(target_warns), 1) | ||
| 559 | + warn_msg = str(target_warns[0].message) | ||
| 560 | + self.assertIn("TorchScript will now drop the function", warn_msg) | ||
| 561 | + | ||
| 562 | + x = torch.tensor(2.0).to(device_type) | ||
| 563 | + eager_out = m(x) | ||
| 564 | + self.assertEqual(eager_out, x) | ||
| 565 | + | ||
| 566 | + m.ignored_func(x) | ||
| 567 | + self.assertEqual(m.state, torch.tensor([999.0], device=device_type)) | ||
| 568 | + self.assertEqual(m.ignored_func(torch.tensor(3, device=device_type)), torch.tensor(30, device=device_type)) | ||
| 569 | + | ||
| 570 | + m_export = self.getExportImportCopy(m) | ||
| 571 | + | ||
| 572 | + with self.assertRaises(AttributeError): | ||
| 573 | + _ = m_export.ignored_func | ||
| 574 | + | ||
| 575 | + self.assertTrue(hasattr(m, "ignored_func")) | ||
| 576 | + self.assertFalse(hasattr(m_export, "ignored_func")) | ||
| 577 | + self.assertNotIn("ignored_func", dir(m_export)) | ||
| 578 | + | ||
| 579 | + export_out = m_export(x) | ||
| 580 | + self.assertEqual(export_out, x) | ||
| 581 | + self.assertEqual(m_export.state, torch.tensor([999.0], device=device_type)) | ||
| 582 | + | ||
| 583 | + def test_ignored_props(self): | ||
| 584 | + class A(nn.Module): | ||
| 585 | + __jit_ignored_attributes__ = ["ignored", "ignored_return_val"] | ||
| 586 | + | ||
| 587 | + | ||
| 588 | + def ignored(self): | ||
| 589 | + raise ValueError("shouldn't be called") | ||
| 590 | + | ||
| 591 | + | ||
| 592 | + def ignored_return_val(self): | ||
| 593 | + return 1 | ||
| 594 | + | ||
| 595 | + | ||
| 596 | + def call(self): | ||
| 597 | + return self.ignored_return_val | ||
| 598 | + | ||
| 599 | + f = torch.jit.script(A()) | ||
| 600 | + # jank way to test if there is no error | ||
| 601 | + self.assertTrue(isinstance(f, torch.jit.ScriptModule)) | ||
| 602 | + self.assertTrue(isinstance(f.call(), property)) | ||
| 603 | + | ||
| 604 | + def test_torch_ignore_conversion_to_none(self): | ||
| 605 | + class A(torch.nn.Module): | ||
| 606 | + | ||
| 607 | + def ignored(self, a: int) -> None: | ||
| 608 | + l: int = len([2 for i in range(a) if i > 2]) | ||
| 609 | + return | ||
| 610 | + | ||
| 611 | + def forward(self) -> int: | ||
| 612 | + a: int = 4 | ||
| 613 | + b: int = 5 | ||
| 614 | + self.ignored(a) | ||
| 615 | + return a + b | ||
| 616 | + | ||
| 617 | + class B(torch.nn.Module): | ||
| 618 | + | ||
| 619 | + def ignored(self, a: int): | ||
| 620 | + l: int = len([2 for i in range(a) if i > 2]) | ||
| 621 | + return | ||
| 622 | + | ||
| 623 | + def forward(self) -> int: | ||
| 624 | + a: int = 4 | ||
| 625 | + b: int = 5 | ||
| 626 | + self.ignored(a) | ||
| 627 | + return a + b | ||
| 628 | + | ||
| 629 | + modelA = torch.jit.script(A()) | ||
| 630 | + self.assertEqual(modelA(), 9) | ||
| 631 | + | ||
| 632 | + modelB = torch.jit.script(B()) | ||
| 633 | + self.assertEqual(modelB(), 9) | ||
| 634 | + | ||
| 635 | + def test_comment_ignore_indent(self): | ||
| 636 | + class Model(torch.nn.Module): | ||
| 637 | + def __init__(self) -> None: | ||
| 638 | + # useless comment that is not indented correctly # noqa: E115 | ||
| 639 | + super().__init__() | ||
| 640 | + | ||
| 641 | + def forward(self): | ||
| 642 | + return 5 | ||
| 643 | + | ||
| 644 | + # should compile without an error | ||
| 645 | + self.checkModule(Model(), ()) | ||
| 646 | + | ||
| 647 | + def test_ignored_method_binding(self): | ||
| 648 | + class Bar(torch.nn.Module): | ||
| 649 | + def __init__(self) -> None: | ||
| 650 | + super().__init__() | ||
| 651 | + self.x : int = 0 | ||
| 652 | + | ||
| 653 | + | ||
| 654 | + def setx(self, x : int): | ||
| 655 | + self.x = x | ||
| 656 | + | ||
| 657 | + | ||
| 658 | + def getx(self): | ||
| 659 | + return self.x | ||
| 660 | + | ||
| 661 | + | ||
| 662 | + def ignored_getx(self): | ||
| 663 | + return self.x | ||
| 664 | + | ||
| 665 | + b = Bar() | ||
| 666 | + b.setx(123) | ||
| 667 | + sb = torch.jit.script(b) | ||
| 668 | + self.assertEqual(sb.getx(), 123) | ||
| 669 | + self.assertEqual(sb.ignored_getx(), 123) | ||
| 670 | + | ||
| 671 | + sb.setx(456) | ||
| 672 | + self.assertEqual(sb.getx(), 456) | ||
| 673 | + self.assertEqual(sb.ignored_getx(), 456) | ||
| 674 | + | ||
| 675 | + def test_no_self_arg_ignore_function(self): | ||
| 676 | + class MyModule(nn.Module): | ||
| 677 | + | ||
| 678 | + def call_np(): | ||
| 679 | + return np.random.choice(2, p=[.95, .05]) | ||
| 680 | + | ||
| 681 | + def forward(self): | ||
| 682 | + return self.call_np() | ||
| 683 | + | ||
| 684 | + with self.assertRaisesRegex(Exception, "does not have a self argument"): | ||
| 685 | + torch.jit.script(MyModule()) | ||
| 686 | + | ||
| 525 | if __name__ == "__main__": | 687 | if __name__ == "__main__": |
| 526 | run_tests() | 688 | run_tests() |