Wrap a DistributedDataParallel module and forward requests for missing attributes to the module wrapped by DDP (the twice-wrapped module). Also forward calls to :func:`state_dict` and :func:`load_state_dict`. Usage:: module.xyz = "hello world" wrapped_module = Dist
| 7 | |
| 8 | |
| 9 | class ModuleProxyWrapper(nn.Module): |
| 10 | """ |
| 11 | Wrap a DistributedDataParallel module and forward requests for missing |
| 12 | attributes to the module wrapped by DDP (the twice-wrapped module). |
| 13 | Also forward calls to :func:`state_dict` and :func:`load_state_dict`. |
| 14 | |
| 15 | Usage:: |
| 16 | |
| 17 | module.xyz = "hello world" |
| 18 | wrapped_module = DistributedDataParallel(module, **ddp_args) |
| 19 | wrapped_module = ModuleProxyWrapper(wrapped_module) |
| 20 | assert wrapped_module.xyz == "hello world" |
| 21 | assert wrapped_module.state_dict().keys() == module.state_dict().keys() |
| 22 | |
| 23 | Args: |
| 24 | module (nn.Module): module to wrap |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, module: nn.Module): |
| 28 | super().__init__() |
| 29 | assert hasattr(module, "module"), \ |
| 30 | "ModuleProxyWrapper expects input to wrap another module" |
| 31 | self.module = module |
| 32 | |
| 33 | def __getattr__(self, name): |
| 34 | """Forward missing attributes to twice-wrapped module.""" |
| 35 | try: |
| 36 | # defer to nn.Module's logic |
| 37 | return super().__getattr__(name) |
| 38 | except AttributeError: |
| 39 | try: |
| 40 | # forward to the once-wrapped module |
| 41 | return getattr(self.module, name) |
| 42 | except AttributeError: |
| 43 | # forward to the twice-wrapped module |
| 44 | return getattr(self.module.module, name) |
| 45 | |
| 46 | def state_dict(self, *args, **kwargs): |
| 47 | """Forward to the twice-wrapped module.""" |
| 48 | return self.module.module.state_dict(*args, **kwargs) |
| 49 | |
| 50 | def load_state_dict(self, *args, **kwargs): |
| 51 | """Forward to the twice-wrapped module.""" |
| 52 | return self.module.module.load_state_dict(*args, **kwargs) |
| 53 | |
| 54 | def forward(self, *args, **kwargs): |
| 55 | return self.module(*args, **kwargs) |
no outgoing calls
no test coverage detected