| 86 | |
| 87 | |
| 88 | class BaseModule(nn.Module, Updateable): |
| 89 | @dataclass |
| 90 | class Config: |
| 91 | weights: Optional[str] = None |
| 92 | freeze: Optional[bool] = False |
| 93 | |
| 94 | cfg: Config # add this to every subclass of BaseModule to enable static type checking |
| 95 | |
| 96 | def __init__( |
| 97 | self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs |
| 98 | ) -> None: |
| 99 | super().__init__() |
| 100 | self.cfg = parse_structured(self.Config, cfg) |
| 101 | self.device = get_device() |
| 102 | self._non_modules = {} |
| 103 | self.configure(*args, **kwargs) |
| 104 | if self.cfg.weights is not None: |
| 105 | # format: path/to/weights:module_name |
| 106 | weights_path, module_name = self.cfg.weights.split(":") |
| 107 | state_dict = load_module_weights( |
| 108 | weights_path, module_name=module_name, map_location="cpu" |
| 109 | ) |
| 110 | self.load_state_dict(state_dict, strict=False) |
| 111 | # self.do_update_step( |
| 112 | # epoch, global_step, on_load_weights=True |
| 113 | # ) # restore states |
| 114 | |
| 115 | if self.cfg.freeze: |
| 116 | for params in self.parameters(): |
| 117 | params.requires_grad = False |
| 118 | |
| 119 | def configure(self, *args, **kwargs) -> None: |
| 120 | pass |
| 121 | |
| 122 | def register_non_module(self, name: str, module: nn.Module) -> None: |
| 123 | # non-modules won't be treated as model parameters |
| 124 | if name in self._non_modules: |
| 125 | raise ValueError(f"Non-module {name} already exists!") |
| 126 | self._non_modules[name] = module |
| 127 | |
| 128 | def non_module(self, name: str): |
| 129 | return self._non_modules.get(name, None) |
nothing calls this directly
no outgoing calls
no test coverage detected