| 160 | |
| 161 | |
| 162 | def load_model(model, model_file, is_restore=False): |
| 163 | t_start = time.time() |
| 164 | |
| 165 | if model_file is None: |
| 166 | return model |
| 167 | |
| 168 | if isinstance(model_file, str): |
| 169 | state_dict = torch.load(model_file) |
| 170 | if "model" in state_dict.keys(): |
| 171 | state_dict = state_dict["model"] |
| 172 | elif "state_dict" in state_dict.keys(): |
| 173 | state_dict = state_dict["state_dict"] |
| 174 | elif "module" in state_dict.keys(): |
| 175 | state_dict = state_dict["module"] |
| 176 | else: |
| 177 | state_dict = model_file |
| 178 | t_ioend = time.time() |
| 179 | |
| 180 | if is_restore: |
| 181 | new_state_dict = OrderedDict() |
| 182 | for k, v in state_dict.items(): |
| 183 | name = "module." + k |
| 184 | new_state_dict[name] = v |
| 185 | state_dict = new_state_dict |
| 186 | |
| 187 | model.load_state_dict(state_dict, strict=True) |
| 188 | ckpt_keys = set(state_dict.keys()) |
| 189 | own_keys = set(model.state_dict().keys()) |
| 190 | missing_keys = own_keys - ckpt_keys |
| 191 | unexpected_keys = ckpt_keys - own_keys |
| 192 | |
| 193 | del state_dict |
| 194 | t_end = time.time() |
| 195 | logger.info( |
| 196 | "Load model, Time usage:\n\tIO: {}, initialize parameters: {}".format(t_ioend - t_start, t_end - t_ioend) |
| 197 | ) |
| 198 | |
| 199 | return model |
| 200 | |
| 201 | |
| 202 | def parse_devices(input_devices): |