Save checkpoint to file. The checkpoint will have 3 fields: ``meta``, ``state_dict`` and ``optimizer``. By default, ``meta`` will contain version and time info. Args: model (Module): Module whose params are to be saved. filename (str): Checkpoint filename. optim
(model: torch.nn.Module,
filename: str,
optimizer: Optional[Optimizer] = None,
lr_scheduler: Optional[_LRScheduler] = None,
meta: Optional[dict] = None,
with_meta: bool = True,
with_model: bool = True)
| 45 | |
| 46 | |
| 47 | def save_checkpoint(model: torch.nn.Module, |
| 48 | filename: str, |
| 49 | optimizer: Optional[Optimizer] = None, |
| 50 | lr_scheduler: Optional[_LRScheduler] = None, |
| 51 | meta: Optional[dict] = None, |
| 52 | with_meta: bool = True, |
| 53 | with_model: bool = True) -> None: |
| 54 | """Save checkpoint to file. |
| 55 | |
| 56 | The checkpoint will have 3 fields: ``meta``, ``state_dict`` and |
| 57 | ``optimizer``. By default, ``meta`` will contain version and time info. |
| 58 | |
| 59 | Args: |
| 60 | model (Module): Module whose params are to be saved. |
| 61 | filename (str): Checkpoint filename. |
| 62 | optimizer (:obj:`Optimizer`, optional): Optimizer to be saved. |
| 63 | lr_scheduler(:obj:`_LRScheduler`, optional): LRScheduler to be saved. |
| 64 | meta (dict, optional): Metadata to be saved in checkpoint. |
| 65 | with_meta (bool, optional): Save meta info. |
| 66 | with_model(bool, optional): Save model states. |
| 67 | """ |
| 68 | checkpoint = {} |
| 69 | if not with_meta and not with_model: |
| 70 | raise ValueError( |
| 71 | 'Save meta by "with_meta=True" or model by "with_model=True"') |
| 72 | |
| 73 | if with_meta: |
| 74 | if meta is None: |
| 75 | meta = {} |
| 76 | elif not isinstance(meta, dict): |
| 77 | raise TypeError( |
| 78 | f'meta must be a dict or None, but got {type(meta)}') |
| 79 | from modelscope import __version__ |
| 80 | meta.update(modelscope=__version__, time=time.asctime()) |
| 81 | |
| 82 | if isinstance(model, torch.nn.parallel.DistributedDataParallel): |
| 83 | model = model.module |
| 84 | |
| 85 | if hasattr(model, 'CLASSES') and model.CLASSES is not None: |
| 86 | # save class name to the meta |
| 87 | meta.update(CLASSES=model.CLASSES) |
| 88 | |
| 89 | checkpoint['meta'] = meta |
| 90 | |
| 91 | # save optimizer state dict in the checkpoint |
| 92 | if isinstance(optimizer, Optimizer): |
| 93 | checkpoint['optimizer'] = optimizer.state_dict() |
| 94 | elif isinstance(optimizer, dict): |
| 95 | checkpoint['optimizer'] = {} |
| 96 | for name, optim in optimizer.items(): |
| 97 | checkpoint['optimizer'][name] = optim.state_dict() |
| 98 | |
| 99 | # save lr_scheduler state dict in the checkpoint |
| 100 | if lr_scheduler is not None and hasattr(lr_scheduler, 'state_dict'): |
| 101 | checkpoint['lr_scheduler'] = lr_scheduler.state_dict() |
| 102 | |
| 103 | if with_model: |
| 104 | if isinstance(model, torch.nn.parallel.DistributedDataParallel): |
no test coverage detected
searching dependent graphs…