The base model. It is a convenient class the collects necessary information (e.g., normalization statistics, codemap, etc) with the model so they can be tracked in the future. You should define two functions: - forward: this is called during `training` (via Model(
| 19 | |
| 20 | |
| 21 | class BaseModel: |
| 22 | """ |
| 23 | The base model. It is a convenient class the collects |
| 24 | necessary information (e.g., normalization statistics, |
| 25 | codemap, etc) with the model so they can be tracked in |
| 26 | the future. |
| 27 | |
| 28 | You should define two functions: |
| 29 | |
| 30 | - forward: |
| 31 | this is called during `training` (via Model(...)) |
| 32 | - infer: |
| 33 | this is called during `inference` (via Model.infer(...)) |
| 34 | |
| 35 | Note that it is not a `nn.Module`. It defines a collection |
| 36 | of `nn.Modules` and their interaction during training and inference. |
| 37 | |
| 38 | Instead of defining as a single nn.Module, this way we allows more easily |
| 39 | using different optimizer for different sub-nn.Module. |
| 40 | """ |
| 41 | |
| 42 | def __init__(self): |
| 43 | super().__init__() |
| 44 | self.var_names_to_save: T.Set[str] = set() |
| 45 | self.var_names_to_load: T.Set[str] = set() |
| 46 | self.buffer_names: T.Set[str] = set() |
| 47 | self.parameter_names: T.Set[str] = set() |
| 48 | self.device = torch.device("cpu") |
| 49 | |
| 50 | def register_buffer(self, name: str, tensor: torch.Tensor, persistent=True): |
| 51 | self.buffer_names.add(name) |
| 52 | setattr(self, name, tensor) |
| 53 | |
| 54 | if persistent: |
| 55 | self.register_var_to_load(var_name=name) |
| 56 | else: |
| 57 | self.register_var_to_save(var_name=name) |
| 58 | |
| 59 | def register_parameter(self, name: str, param: torch.Tensor): |
| 60 | self.parameter_names.add(name) |
| 61 | setattr(self, name, param) |
| 62 | self.register_var_to_load(var_name=name) |
| 63 | |
| 64 | def register_var_to_save( |
| 65 | self, |
| 66 | var_name: T.Union[str, T.List[str]], |
| 67 | ): |
| 68 | """ |
| 69 | Make sure the var will be saved in the state_dict. |
| 70 | Note that it is just for recording purpose. |
| 71 | They will NOT be loaded during retraining. |
| 72 | """ |
| 73 | # assert torch.__version__ >= "1.10.0" |
| 74 | if isinstance(var_name, str): |
| 75 | var_name = [var_name] |
| 76 | |
| 77 | if getattr(self, "var_names_to_save", None) is None: |
| 78 | self.var_names_to_save = set() |
nothing calls this directly
no outgoing calls
no test coverage detected