(self, device, dtype, module_info, training)
| 513 | torch.float64: tol(4e-4, 0)}) |
| 514 | @modules(module_db) |
| 515 | def test_cpu_gpu_parity(self, device, dtype, module_info, training): |
| 516 | # TODO: RNN / GRU / LSTM don't support backwards on eval mode for cuDNN; skip this in a |
| 517 | # nicer way for eval mode only. |
| 518 | # See https://github.com/pytorch/pytorch/issues/79161 |
| 519 | rnn_modules = {torch.nn.RNN, torch.nn.GRU, torch.nn.LSTM} |
| 520 | if (module_info.module_cls in rnn_modules |
| 521 | and not training |
| 522 | and 'cuda' in device |
| 523 | and torch.backends.cudnn.enabled): |
| 524 | return |
| 525 | |
| 526 | # Test cpu and gpu results are the same |
| 527 | module_cls = module_info.module_cls |
| 528 | module_inputs_cpu = module_info.module_inputs_func(module_info, device="cpu", dtype=dtype, |
| 529 | requires_grad=True, training=training) |
| 530 | |
| 531 | def _to_device(obj): |
| 532 | if isinstance(obj, torch.Tensor): |
| 533 | res = obj.detach().to(device=device) |
| 534 | res.requires_grad = obj.requires_grad |
| 535 | return res |
| 536 | elif isinstance(obj, tuple): |
| 537 | return tuple(_to_device(o) for o in obj) |
| 538 | elif isinstance(obj, dict): |
| 539 | return {key: _to_device(o) for key, o in obj.items()} |
| 540 | else: |
| 541 | return deepcopy(obj) |
| 542 | |
| 543 | for module_input in module_inputs_cpu: |
| 544 | # === Move input from cpu to device === |
| 545 | cpu_forward_args = module_input.forward_input.args |
| 546 | cpu_forward_kwargs = module_input.forward_input.kwargs |
| 547 | |
| 548 | gpu_forward_args, gpu_forward_kwargs = _to_device((cpu_forward_args, cpu_forward_kwargs)) |
| 549 | |
| 550 | self._retain_grad((cpu_forward_args, cpu_forward_kwargs, gpu_forward_args, gpu_forward_kwargs)) |
| 551 | |
| 552 | # === Construct module on cpu and gpu === |
| 553 | args, kwargs = module_input.constructor_input.args, module_input.constructor_input.kwargs |
| 554 | |
| 555 | cpu_module = module_cls(*args, **kwargs).to(dtype).to("cpu") |
| 556 | cpu_module.train(training) |
| 557 | gpu_module = module_cls(*args, **kwargs).to(dtype).to(device) |
| 558 | gpu_module.train(training) |
| 559 | |
| 560 | # === Lazy modules need to see an input to initialize params === |
| 561 | if issubclass(module_cls, torch.nn.modules.lazy.LazyModuleMixin): |
| 562 | with torch.no_grad(): |
| 563 | cpu_module(*cpu_forward_args, **cpu_forward_kwargs) |
| 564 | gpu_module(*gpu_forward_args, **gpu_forward_kwargs) |
| 565 | |
| 566 | for cpu_p, gpu_p in zip(cpu_module.parameters(), gpu_module.parameters()): |
| 567 | gpu_p.data.copy_(cpu_p) |
| 568 | |
| 569 | # === Compare forward output between cpu and gpu === |
| 570 | cpu_outputs = cpu_module(*cpu_forward_args, **cpu_forward_kwargs) |
| 571 | gpu_outputs = gpu_module(*gpu_forward_args, **gpu_forward_kwargs) |
| 572 |
nothing calls this directly
no test coverage detected