File size: 1,687 Bytes
b84549f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from typing import Dict, Any
from torch import nn
from data.datasets.ab_dataset import ABDataset
from abc import ABC, abstractmethod
from utils.common.log import logger
import json
import os
from utils.common.others import backup_key_codes
from .model import BaseModel
from data import Scenario
from schema import Schema
from utils.common.data_record import write_json
class BaseAlg(ABC):
def __init__(self, models: Dict[str, BaseModel], res_save_dir):
self.models = models
self.res_save_dir = res_save_dir
self.get_required_models_schema().validate(models)
os.makedirs(res_save_dir)
logger.info(f'[alg] init alg: {self.__class__.__name__}, res saved in {res_save_dir}')
@abstractmethod
def get_required_models_schema(self) -> Schema:
raise NotImplementedError
@abstractmethod
def get_required_hyp_schema(self) -> Schema:
raise NotImplementedError
@abstractmethod
def run(self,
scenario: Scenario,
hyps: Dict) -> Dict[str, Any]:
"""
return metrics
"""
self.get_required_hyp_schema().validate(hyps)
try:
write_json(os.path.join(self.res_save_dir, 'hyps.json'), hyps, ensure_obj_serializable=True)
except:
with open(os.path.join(self.res_save_dir, 'hyps.txt'), 'w') as f:
f.write(str(hyps))
write_json(os.path.join(self.res_save_dir, 'scenario.json'), scenario.to_json())
logger.info(f'[alg] alg {self.__class__.__name__} start running')
backup_key_codes(os.path.join(self.res_save_dir, 'backup_codes')) |