(state_dict: Dict[str, Any], prefix: str)
| 9 | |
| 10 | |
| 11 | def append_prefix(state_dict: Dict[str, Any], prefix: str) -> None: |
| 12 | keys = sorted(state_dict.keys()) |
| 13 | if not all(len(key) == 0 or not key.startswith(prefix) for key in keys): |
| 14 | return |
| 15 | |
| 16 | for key in keys: |
| 17 | newkey = prefix + key |
| 18 | state_dict[newkey] = state_dict.pop(key) |
| 19 | |
| 20 | # also strip the prefix in metadata, if any.. |
| 21 | try: |
| 22 | metadata = state_dict._metadata # pyre-ignore |
| 23 | except AttributeError: |
| 24 | pass |
| 25 | else: |
| 26 | for key in list(metadata.keys()): |
| 27 | # for the metadata dict, the key can be: |
| 28 | # '': for the DDP module, which we want to remove. |
| 29 | # 'module': for the actual model. |
| 30 | # 'module.xx.xx': for the rest. |
| 31 | |
| 32 | if len(key) == 0: |
| 33 | continue |
| 34 | newkey = prefix + key |
| 35 | metadata[newkey] = metadata.pop(key) |
| 36 | |
| 37 | class SGTCheckPointer(Checkpointer): |
| 38 | def __init__(self, model, save_dir="", *, save_to_disk=None, **checkpointables): |
nothing calls this directly
no outgoing calls
no test coverage detected