|
import torch |
|
from abc import ABC, abstractmethod |
|
from utils.common.file import ensure_dir |
|
from utils.common.log import logger |
|
import time |
|
from typing import List |
|
|
|
|
|
class BaseModel(ABC): |
|
def __init__(self, |
|
name: str, |
|
models_dict_path: str, |
|
device: str): |
|
|
|
self.name = name |
|
self.models_dict_path = models_dict_path |
|
self.models_dict = torch.load(models_dict_path, map_location=device) |
|
|
|
self.device = device |
|
|
|
assert set(self.get_required_model_components()) <= set(list(self.models_dict.keys())) |
|
|
|
self.to(device) |
|
|
|
logger.info(f'[model] init model: {dict(name=name, components=self.get_required_model_components())}') |
|
logger.debug(self.models_dict) |
|
|
|
@abstractmethod |
|
def get_required_model_components(self) -> List[str]: |
|
pass |
|
|
|
@abstractmethod |
|
def get_accuracy(self, test_loader, *args, **kwargs): |
|
pass |
|
|
|
@abstractmethod |
|
def infer(self, x, *args, **kwargs): |
|
pass |
|
|
|
def save_model(self, p: str): |
|
logger.debug(f'[model] save model: {self.name}') |
|
ensure_dir(p) |
|
torch.save(self.models_dict, p) |
|
|
|
def load_model(self, p: str): |
|
logger.debug(f'[model] load model: {self.name}, from {p}') |
|
self.models_dict = torch.load(p, map_location=self.device) |
|
|
|
def to(self, device): |
|
logger.debug(f'[model] to device: {device}') |
|
for k, v in self.models_dict.items(): |
|
try: |
|
self.models_dict[k] = v.to(device) |
|
except Exception as e: |
|
pass |
|
|
|
def to_eval_mode(self, verbose=False): |
|
if verbose: |
|
logger.info(f'[model] to eval mode') |
|
for k, v in self.models_dict.items(): |
|
try: |
|
self.models_dict[k].eval() |
|
except Exception as e: |
|
pass |
|
|
|
def to_train_mode(self, verbose=False): |
|
if verbose: |
|
logger.info(f'[model] to train mode') |
|
for k, v in self.models_dict.items(): |
|
try: |
|
self.models_dict[k].train() |
|
except Exception as e: |
|
pass |
|
|
