(
model,
state_dict,
start_prefix,
head_prefix_keys,
load_state_fn=None,
)
| 222 | return mismatched_key |
| 223 | |
| 224 | def _load_state_dict_into_model( |
| 225 | model, |
| 226 | state_dict, |
| 227 | start_prefix, |
| 228 | head_prefix_keys, |
| 229 | load_state_fn=None, |
| 230 | ): |
| 231 | # Convert old format to new format if needed from a PyTorch state_dict |
| 232 | old_keys = [] |
| 233 | new_keys = [] |
| 234 | for key in state_dict.keys(): |
| 235 | new_key = None |
| 236 | if 'gamma' in key: |
| 237 | new_key = key.replace('gamma', 'weight') |
| 238 | if 'beta' in key: |
| 239 | new_key = key.replace('beta', 'bias') |
| 240 | if new_key: |
| 241 | old_keys.append(key) |
| 242 | new_keys.append(new_key) |
| 243 | for old_key, new_key in zip(old_keys, new_keys): |
| 244 | state_dict[new_key] = state_dict.pop(old_key) |
| 245 | |
| 246 | # copy state_dict so _load_from_state_dict can modify it |
| 247 | metadata = getattr(state_dict, '_metadata', None) |
| 248 | state_dict = state_dict.copy() |
| 249 | if metadata is not None: |
| 250 | state_dict._metadata = metadata |
| 251 | |
| 252 | error_msgs = [] |
| 253 | |
| 254 | if load_state_fn is not None: |
| 255 | load_state_fn( |
| 256 | model, |
| 257 | state_dict, |
| 258 | prefix=start_prefix, |
| 259 | head_prefix_keys=head_prefix_keys, |
| 260 | local_metadata=None, |
| 261 | error_msgs=error_msgs) |
| 262 | else: |
| 263 | |
| 264 | def load(module: nn.Module, prefix=''): |
| 265 | local_metadata = {} if metadata is None else metadata.get( |
| 266 | prefix[:-1], {}) |
| 267 | args = (state_dict, prefix, local_metadata, True, [], [], |
| 268 | error_msgs) |
| 269 | module._load_from_state_dict(*args) |
| 270 | for name, child in module._modules.items(): |
| 271 | if child is not None: |
| 272 | load(child, prefix + name + '.') |
| 273 | |
| 274 | load(model, prefix=start_prefix) |
| 275 | |
| 276 | return error_msgs |
| 277 | |
| 278 | def _load_checkpoint( |
| 279 | model, |
no test coverage detected
searching dependent graphs…