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