| |
| from abc import ABCMeta, abstractmethod, abstractproperty |
|
|
| import torch |
|
|
|
|
| class PytorchModuleHook(metaclass=ABCMeta): |
| """Base class for PyTorch module hook registers. |
| |
| An instance of a subclass of PytorchModuleHook can be used to |
| register hook to a pytorch module using the `register` method like: |
| hook_register.register(module) |
| |
| Subclasses should add/overwrite the following methods: |
| - __init__ |
| - hook |
| - hook_type |
| """ |
|
|
| @abstractmethod |
| def hook(self, *args, **kwargs): |
| """Hook function.""" |
|
|
| @abstractproperty |
| def hook_type(self) -> str: |
| """Hook type Subclasses should overwrite this function to return a |
| string value in. |
| |
| {`forward`, `forward_pre`, `backward`} |
| """ |
|
|
| def register(self, module): |
| """Register the hook function to the module. |
| |
| Args: |
| module (pytorch module): the module to register the hook. |
| |
| Returns: |
| handle (torch.utils.hooks.RemovableHandle): a handle to remove |
| the hook by calling handle.remove() |
| """ |
| assert isinstance(module, torch.nn.Module) |
|
|
| if self.hook_type == 'forward': |
| h = module.register_forward_hook(self.hook) |
| elif self.hook_type == 'forward_pre': |
| h = module.register_forward_pre_hook(self.hook) |
| elif self.hook_type == 'backward': |
| h = module.register_backward_hook(self.hook) |
| else: |
| raise ValueError(f'Invalid hook type {self.hook}') |
|
|
| return h |
|
|
|
|
| class WeightNormClipHook(PytorchModuleHook): |
| """Apply weight norm clip regularization. |
| |
| The module's parameter will be clip to a given maximum norm before each |
| forward pass. |
| |
| Args: |
| max_norm (float): The maximum norm of the parameter. |
| module_param_names (str|list): The parameter name (or name list) to |
| apply weight norm clip. |
| """ |
|
|
| def __init__(self, max_norm=1.0, module_param_names='weight'): |
| self.module_param_names = module_param_names if isinstance( |
| module_param_names, list) else [module_param_names] |
| self.max_norm = max_norm |
|
|
| @property |
| def hook_type(self): |
| return 'forward_pre' |
|
|
| def hook(self, module, _input): |
| for name in self.module_param_names: |
| assert name in module._parameters, f'{name} is not a parameter' \ |
| f' of the module {type(module)}' |
| param = module._parameters[name] |
|
|
| with torch.no_grad(): |
| m = param.norm().item() |
| if m > self.max_norm: |
| param.mul_(self.max_norm / (m + 1e-6)) |
|
|
