Base class for PyTorch meta-learning modules. These modules accept an additional argument `params` in their `forward` method. Notes ----- Objects inherited from `MetaModule` are fully compatible with PyTorch modules from `torch.nn.Module`. The argument `params` is a diction
| 80 | |
| 81 | |
| 82 | class MetaModule(nn.Module): |
| 83 | """ |
| 84 | Base class for PyTorch meta-learning modules. These modules accept an |
| 85 | additional argument `params` in their `forward` method. |
| 86 | |
| 87 | Notes |
| 88 | ----- |
| 89 | Objects inherited from `MetaModule` are fully compatible with PyTorch |
| 90 | modules from `torch.nn.Module`. The argument `params` is a dictionary of |
| 91 | tensors, with full support of the computation graph (for differentiation). |
| 92 | |
| 93 | Based on SIREN's torchmeta with some additional features/changes. |
| 94 | |
| 95 | All meta weights must not have the batch dimension, as they are later tiled |
| 96 | to the given batch size after unsqueezing the first dimension (e.g. a |
| 97 | weight of dimension [d_out x d_in] is tiled to have the dimension [batch x |
| 98 | d_out x d_in]). Requiring all meta weights to have a batch dimension of 1 |
| 99 | (e.g. [1 x d_out x d_in] from the earlier example) could be a more natural |
| 100 | choice, but this results in silent failures. |
| 101 | """ |
| 102 | |
| 103 | def __init__(self, *args, **kwargs): |
| 104 | super().__init__(*args, **kwargs) |
| 105 | self._meta_state_dict = set() |
| 106 | self._meta_params = set() |
| 107 | |
| 108 | def register_meta_buffer(self, name: str, param: nn.Parameter): |
| 109 | """ |
| 110 | Registers a trainable or nontrainable parameter as a meta buffer. This |
| 111 | can be later retrieved by meta_state_dict |
| 112 | """ |
| 113 | self.register_buffer(name, param) |
| 114 | self._meta_state_dict.add(name) |
| 115 | |
| 116 | def register_meta_parameter(self, name: str, parameter: nn.Parameter): |
| 117 | """ |
| 118 | Registers a meta parameter so it is included in named_meta_parameters |
| 119 | and meta_state_dict. |
| 120 | """ |
| 121 | self.register_parameter(name, parameter) |
| 122 | self._meta_params.add(name) |
| 123 | self._meta_state_dict.add(name) |
| 124 | |
| 125 | def register_meta(self, name: str, parameter: nn.Parameter, trainable: bool = True): |
| 126 | if trainable: |
| 127 | self.register_meta_parameter(name, parameter) |
| 128 | else: |
| 129 | self.register_meta_buffer(name, parameter) |
| 130 | |
| 131 | def register(self, name: str, parameter: nn.Parameter, meta: bool, trainable: bool = True): |
| 132 | if meta: |
| 133 | if trainable: |
| 134 | self.register_meta_parameter(name, parameter) |
| 135 | else: |
| 136 | self.register_meta_buffer(name, parameter) |
| 137 | else: |
| 138 | if trainable: |
| 139 | self.register_parameter(name, parameter) |
nothing calls this directly
no outgoing calls
no test coverage detected