| | |
| |
|
| | |
| | |
| | |
| |
|
| | """ |
| | 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 {} |
| |
|
