Save a checkpoint. Args: model (model): model to save the weight to the checkpoint. optimizer (optim): optimizer to save the historical state. loss_scaler (scaler): scaler for loss. epoch (int): current number of epoch of the model. cfg (CfgNode): con
(path_to_job, model, optimizer, loss_scaler, epoch, cfg)
| 116 | |
| 117 | |
| 118 | def save_checkpoint(path_to_job, model, optimizer, loss_scaler, epoch, cfg): |
| 119 | """ |
| 120 | Save a checkpoint. |
| 121 | Args: |
| 122 | model (model): model to save the weight to the checkpoint. |
| 123 | optimizer (optim): optimizer to save the historical state. |
| 124 | loss_scaler (scaler): scaler for loss. |
| 125 | epoch (int): current number of epoch of the model. |
| 126 | cfg (CfgNode): configs to save. |
| 127 | """ |
| 128 | # Save checkpoints only from the master process. |
| 129 | if not du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS): |
| 130 | return |
| 131 | # Ensure that the checkpoint dir exists. |
| 132 | g_pathmgr.mkdirs(get_checkpoint_dir(path_to_job)) |
| 133 | # Omit the DDP wrapper in the multi-gpu setting. |
| 134 | sd = model.module.state_dict() if cfg.NUM_GPUS > 1 else model.state_dict() |
| 135 | normalized_sd = sub_to_normal_bn(sd) |
| 136 | |
| 137 | # Record the state. |
| 138 | checkpoint = { |
| 139 | "epoch": epoch, |
| 140 | "model_state": normalized_sd, |
| 141 | "optimizer_state": optimizer.state_dict(), |
| 142 | 'scaler': loss_scaler.state_dict(), |
| 143 | "cfg": cfg.dump(), |
| 144 | } |
| 145 | # Write the checkpoint. |
| 146 | path_to_checkpoint = get_path_to_checkpoint(path_to_job, epoch + 1) |
| 147 | with g_pathmgr.open(path_to_checkpoint, "wb") as f: |
| 148 | torch.save(checkpoint, f) |
| 149 | return path_to_checkpoint |
| 150 | |
| 151 | |
| 152 | def inflate_weight(state_dict_2d, state_dict_3d): |
nothing calls this directly
no test coverage detected