| |
|
|
| |
| |
| |
|
|
| """ |
| Base class for trainable models. |
| """ |
|
|
| from abc import ABCMeta, abstractmethod |
| from copy import copy |
|
|
| from omegaconf import OmegaConf |
| from torch import nn |
|
|
|
|
| class BaseModel(nn.Module, metaclass=ABCMeta): |
|
|
| required_data_keys = [] |
| strict_conf = True |
|
|
| def __init__(self, conf): |
| """Perform some logic and call the _init method of the child model.""" |
| super().__init__() |
| self.conf = conf |
| OmegaConf.set_readonly(conf, True) |
| OmegaConf.set_struct(conf, True) |
| self.required_data_keys = copy(self.required_data_keys) |
| self._init(conf) |
|
|
| def forward(self, data): |
| """Check the data and call the _forward method of the child model.""" |
|
|
| def recursive_key_check(expected, given): |
| for key in expected: |
| assert key in given, f"Missing key {key} in data" |
| if isinstance(expected, dict): |
| recursive_key_check(expected[key], given[key]) |
|
|
| recursive_key_check(self.required_data_keys, data) |
| return self._forward(data) |
|
|
| @abstractmethod |
| def _init(self, conf): |
| """To be implemented by the child class.""" |
| raise NotImplementedError |
|
|
| @abstractmethod |
| def _forward(self, data): |
| """To be implemented by the child class.""" |
| raise NotImplementedError |
|
|
| def loss(self, pred, data): |
| """To be implemented by the child class.""" |
| raise NotImplementedError |
|
|
| def metrics(self): |
| return {} |
|
|
