Load a model (:py:class:`BaseModel`) to device and as eval mode. Args: filename: filename of the pt file model_names: a list of name of the model in the pth file. In other words, we load `checkpoint[model_name]`. model_classes:
(
filename: str,
model_names: T.Union[str, T.List[str]],
model_classes: T.Union[T.Callable, T.List[T.Callable]],
model_params_names: T.Optional[T.Union[str, T.List[str]]] = None,
model_patch_params: T.Optional[T.Union[T.Dict[str, T.Any], T.List[T.Dict[str, T.Any]]]] = None,
device=torch.device("cpu"),
)
| 12 | |
| 13 | |
| 14 | def load_model( |
| 15 | filename: str, |
| 16 | model_names: T.Union[str, T.List[str]], |
| 17 | model_classes: T.Union[T.Callable, T.List[T.Callable]], |
| 18 | model_params_names: T.Optional[T.Union[str, T.List[str]]] = None, |
| 19 | model_patch_params: T.Optional[T.Union[T.Dict[str, T.Any], T.List[T.Dict[str, T.Any]]]] = None, |
| 20 | device=torch.device("cpu"), |
| 21 | ) -> T.Tuple[T.Dict[str, BaseModel], T.Dict[str, T.Any]]: |
| 22 | """ |
| 23 | Load a model (:py:class:`BaseModel`) to device and as eval mode. |
| 24 | |
| 25 | Args: |
| 26 | filename: |
| 27 | filename of the pt file |
| 28 | model_names: |
| 29 | a list of name of the model in the pth file. |
| 30 | In other words, we load `checkpoint[model_name]`. |
| 31 | model_classes: |
| 32 | a list of class definition of the model class. |
| 33 | The model will load the state_dict from `checkpoint[model_name]`. |
| 34 | model_params_names: |
| 35 | default model params, ie, if the parameters of the model is not |
| 36 | saved in checkpoint[model_name], the model will be initialized |
| 37 | model_params_names. |
| 38 | model_patch_params: |
| 39 | manually set the parameters to overwrite the config_dict |
| 40 | of the model stored in the checkpoint file. |
| 41 | For example, if checkpoint[model_names[0]]['config_dict']['param1'] = 30, |
| 42 | but model_patch_params[0]['param1'] = 0, the model will be created with |
| 43 | `param1 = 0`. |
| 44 | device: |
| 45 | device to load the models. |
| 46 | |
| 47 | Returns: |
| 48 | model_dict: |
| 49 | a dict containing models (model_name -> model) |
| 50 | checkpoint: |
| 51 | checkpoint dict. The dictionary stored by in the pretrained model .pth file. |
| 52 | """ |
| 53 | |
| 54 | assert os.path.exists(filename) |
| 55 | |
| 56 | if isinstance(model_names, str): |
| 57 | model_names = [model_names] |
| 58 | if not isinstance(model_classes, (list, tuple)): |
| 59 | model_classes = [model_classes] |
| 60 | if isinstance(model_params_names, str): |
| 61 | model_params_names = [model_params_names] |
| 62 | if model_params_names is None: |
| 63 | model_params_names = [None] * len(model_names) |
| 64 | assert len(model_names) == len(model_classes) |
| 65 | assert len(model_names) == len(model_params_names) |
| 66 | if model_patch_params is None: |
| 67 | model_patch_params = [dict()] * len(model_names) |
| 68 | if not isinstance(model_patch_params, (list, tuple)): |
| 69 | model_patch_params = [model_patch_params] |
| 70 | |
| 71 | # load the model |