__init__.py

import importlib
import os
import torch
from typing import Tuple, Union
from . import base


# Automatically import any python files this directory
package_dir = os.path.dirname(__file__)
package = os.path.basename(package_dir)
for file in os.listdir(package_dir):
    path = os.path.join(package_dir, file)
    if file.startswith('_') or file.startswith('.'):
        continue
    if file.endswith('.py') or os.path.isdir(path):
        model_name = file[:-3] if file.endswith('.py') else file
        importlib.import_module(f'{package}.{model_name}')


def get_class(model_class_name: str) -> type:
    return base.model_classes[model_class_name]


def create(model_class_name: str, args0: dict, **extra_args) -> base.BaseModel:
    model_class = get_class(model_class_name)
    return model_class(args0, extra_args)


def load(path: Union[str, os.PathLike], args0: dict = {}, **extra_args) -> Tuple[base.BaseModel, dict]:
    states: dict = torch.load(path)
    states['args'].update(args0)
    model = create(states['model'], states['args'], **extra_args)
    model.load_state_dict(states['states'])
    return model, states


def save(path: Union[str, os.PathLike], model: base.BaseModel, **extra_states):
    #print(f'Save model to {path}...')
    dict = {
        'model': model.__class__.__name__,
        'args': model.args0,
        'states': model.state_dict(),
        **extra_states
    }
    torch.save(dict, path)