(self, device, dtype, module_info, training)
| 604 | @with_tf32_off |
| 605 | @modules(module_db) |
| 606 | def test_memory_format(self, device, dtype, module_info, training): |
| 607 | is_sm86or80 = device.startswith("cuda") and (torch.cuda.get_device_capability(0) == (8, 6) |
| 608 | or torch.cuda.get_device_capability(0) == (8, 0)) |
| 609 | # TODO tighten it to a specific module |
| 610 | atol, rtol = (3e-3, 7e-3) if is_sm86or80 else (None, None) |
| 611 | module_cls = module_info.module_cls |
| 612 | module_inputs = module_info.module_inputs_func(module_info, device=device, dtype=dtype, |
| 613 | requires_grad=True, training=training) |
| 614 | module_memformat_affects_out = module_info.module_memformat_affects_out |
| 615 | |
| 616 | def _get_mem_formats(channels_last=False, channels_last_3d=False): |
| 617 | if channels_last: |
| 618 | return ([torch.contiguous_format, torch.channels_last], |
| 619 | [torch.preserve_format, torch.contiguous_format, torch.channels_last]) |
| 620 | elif channels_last_3d: |
| 621 | return ([torch.contiguous_format, torch.channels_last_3d], |
| 622 | [torch.preserve_format, torch.contiguous_format, torch.channels_last_3d]) |
| 623 | else: |
| 624 | return ([torch.contiguous_format], |
| 625 | [torch.preserve_format, torch.contiguous_format]) |
| 626 | |
| 627 | # Check that at least one Tensor input has dim == n |
| 628 | def _check_dims(obj, n): |
| 629 | if isinstance(obj, torch.Tensor): |
| 630 | return obj.dim() == n |
| 631 | elif isinstance(obj, (tuple, list)): |
| 632 | return any(_check_dims(o, n) for o in obj) |
| 633 | else: |
| 634 | return False |
| 635 | |
| 636 | # Called after _check_dims, when we know that >= 1 tensor can be converted to mem_format |
| 637 | def _to_mem_format(mem_format, obj): |
| 638 | def inner_to_mem_format(obj): |
| 639 | d = obj.dim() |
| 640 | if ((mem_format == torch.channels_last and d != 4) |
| 641 | or (mem_format == torch.channels_last_3d and d != 5)): |
| 642 | return obj.clone().detach().requires_grad_(obj.requires_grad) |
| 643 | return obj.clone().to(memory_format=mem_format).detach().requires_grad_(obj.requires_grad) |
| 644 | |
| 645 | return self._traverse_obj(obj, inner_to_mem_format) |
| 646 | |
| 647 | def _check_out_mem_format(output, input_mem_format, module_mem_format): |
| 648 | def inner_check_out_mem_format(output): |
| 649 | d = output.dim() |
| 650 | if (d == 4 and ((input_mem_format == torch.channels_last) |
| 651 | or (module_mem_format == torch.channels_last and module_memformat_affects_out))): |
| 652 | self.assertTrue(output.is_contiguous(memory_format=torch.channels_last)) |
| 653 | elif (d == 5 and ((input_mem_format == torch.channels_last_3d) |
| 654 | or (module_mem_format == torch.channels_last_3d and module_memformat_affects_out))): |
| 655 | self.assertTrue(output.is_contiguous(memory_format=torch.channels_last_3d)) |
| 656 | else: |
| 657 | self.assertTrue(output.is_contiguous()) |
| 658 | return self._traverse_obj(output, inner_check_out_mem_format) |
| 659 | |
| 660 | def _req_grad(t): |
| 661 | return isinstance(t, torch.Tensor) and t.requires_grad |
| 662 | |
| 663 | for module_input in module_inputs: |
nothing calls this directly
no test coverage detected