RabbitRUI/ruispace

1"""This package contains modules related to objective functions, optimizations, and network architectures.2 3To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.4You need to implement the following five functions:5    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).6    -- <set_input>:                     unpack data from dataset and apply preprocessing.7    -- <forward>:                       produce intermediate results.8    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.9    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.10 11In the function <__init__>, you need to define four lists:12    -- self.loss_names (str list):          specify the training losses that you want to plot and save.13    -- self.model_names (str list):         define networks used in our training.14    -- self.visual_names (str list):        specify the images that you want to display and save.15    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.16 17Now you can use the model class by specifying flag '--model dummy'.18See our template model class 'template_model.py' for more details.19"""20 21import importlib22from src.face3d.models.base_model import BaseModel23 24 25def find_model_using_name(model_name):26    """Import the module "models/[model_name]_model.py".27 28    In the file, the class called DatasetNameModel() will29    be instantiated. It has to be a subclass of BaseModel,30    and it is case-insensitive.31    """32    model_filename = "face3d.models." + model_name + "_model"33    modellib = importlib.import_module(model_filename)34    model = None35    target_model_name = model_name.replace('_', '') + 'model'36    for name, cls in modellib.__dict__.items():37        if name.lower() == target_model_name.lower() \38           and issubclass(cls, BaseModel):39            model = cls40 41    if model is None:42        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))43        exit(0)44 45    return model46 47 48def get_option_setter(model_name):49    """Return the static method <modify_commandline_options> of the model class."""50    model_class = find_model_using_name(model_name)51    return model_class.modify_commandline_options52 53 54def create_model(opt):55    """Create a model given the option.56 57    This function warps the class CustomDatasetDataLoader.58    This is the main interface between this package and 'train.py'/'test.py'59 60    Example:61        >>> from models import create_model62        >>> model = create_model(opt)63    """64    model = find_model_using_name(opt.model)65    instance = model(opt)66    print("model [%s] was created" % type(instance).__name__)67    return instance68