| 131 | |
| 132 | |
| 133 | def load_restore_model(model, model_file): |
| 134 | t_start = time.time() |
| 135 | |
| 136 | if model_file is None: |
| 137 | return model |
| 138 | |
| 139 | if isinstance(model_file, str): |
| 140 | state_dict = torch.load(model_file) |
| 141 | if "model" in state_dict.keys(): |
| 142 | state_dict = state_dict["model"] |
| 143 | elif "state_dict" in state_dict.keys(): |
| 144 | state_dict = state_dict["state_dict"] |
| 145 | elif "module" in state_dict.keys(): |
| 146 | state_dict = state_dict["module"] |
| 147 | else: |
| 148 | state_dict = model_file |
| 149 | t_ioend = time.time() |
| 150 | |
| 151 | model.load_state_dict(state_dict, strict=True) |
| 152 | |
| 153 | del state_dict |
| 154 | t_end = time.time() |
| 155 | logger.info( |
| 156 | "Load model, Time usage:\n\tIO: {}, initialize parameters: {}".format(t_ioend - t_start, t_end - t_ioend) |
| 157 | ) |
| 158 | |
| 159 | return model |
| 160 | |
| 161 | |
| 162 | def load_model(model, model_file, is_restore=False): |