A TestCase method that traces/exports/tests an nn.Module and method.
(self: unittest.TestCase)
| 494 | """ |
| 495 | |
| 496 | def wrapper(self: unittest.TestCase) -> None: |
| 497 | """A TestCase method that traces/exports/tests an nn.Module and method.""" |
| 498 | module = ExportedModule.export( |
| 499 | module_class=module_cls, |
| 500 | # testend2end only supports modules with single methods defined |
| 501 | methods=(method,), |
| 502 | ignore_to_out_var_failure=ignore_to_out_var_failure, |
| 503 | dynamic_memory_planning_mode=dynamic_memory_planning_mode, |
| 504 | ) |
| 505 | if verify_graph: |
| 506 | verify_graph(self, module.exported_program.graph_module) |
| 507 | print(f"inputs for tracing: {module.trace_inputs}") |
| 508 | |
| 509 | # compare the result between the eager module and graph module |
| 510 | inputs_list = [module.get_random_inputs() for _ in range(niter)] |
| 511 | |
| 512 | if run_graph_module: |
| 513 | for inputs in inputs_list: |
| 514 | with torch.no_grad(): |
| 515 | # only one method is supported so just grab that single method |
| 516 | expected = getattr(module.eager_module, module.methods[0])(*inputs) |
| 517 | with torch.no_grad(): |
| 518 | result = module.exported_program.module()(*inputs) |
| 519 | self.assertTrue(allclose(expected, result, rtol, atol)) |
| 520 | |
| 521 | program = module.executorch_program.executorch_program |
| 522 | pretty_print(program) |
| 523 | print_program(program, show_meminfo=True, mark_dynamic_shape_tensor=True) |
| 524 | print(f"mem buffer sizes: {program.execution_plan[0].non_const_buffer_sizes}") |
| 525 | if not allow_non_contiguous_tensor: |
| 526 | validate_contiguous_tensors(program) |
| 527 | self.assertTrue(len(program.execution_plan[0].non_const_buffer_sizes) >= 2) |
| 528 | # We should not enable the following assertion since for some models |
| 529 | # that simply returning graph input, no mutable memory should be allocated |
| 530 | # self.assertTrue(all(s > 0 for s in program.program.execution_plan[0].non_const_buffer_sizes[1:])) |
| 531 | |
| 532 | program.version = 0 |
| 533 | buff = module.executorch_program.buffer |
| 534 | # Check that the magic version number is in the expected place, and |
| 535 | # follows the expected pattern. |
| 536 | self.assertRegex(buff[4:8].decode(errors="replace"), r"^ET[0-9][0-9]$") |
| 537 | |
| 538 | if run_executor: |
| 539 | print("Running on the runtime") |
| 540 | executorch_module = _load_for_executorch_from_buffer(buff) |
| 541 | # compare the result between eager module and executor |
| 542 | for idx, inputs in enumerate(inputs_list): |
| 543 | with torch.no_grad(): |
| 544 | expected = getattr(module.eager_module, method)(*inputs) |
| 545 | |
| 546 | if do_tree_flatten: |
| 547 | # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. |
| 548 | flatten_inputs, inputs_spec = pytree.tree_flatten(*inputs) |
| 549 | executorch_result = executorch_module.forward([*flatten_inputs]) |
| 550 | # pyre-fixme[16]: Module `pytree` has no attribute `TreeSpec`. |
| 551 | executorch_result_unflatten = pytree.TreeSpec.from_str( |
| 552 | program.execution_plan[0].container_meta_type.encoded_out_str |
| 553 | ).tree_unflatten(executorch_result) |
nothing calls this directly
no test coverage detected