item_dict: {"model": model, "opt1": opt1, ...}
(item_dict, model_dir, map_location=None, step=None)
| 23 | |
| 24 | |
| 25 | def load_checkpoint(item_dict, model_dir, map_location=None, step=None): |
| 26 | """ item_dict: {"model": model, "opt1": opt1, ...}""" |
| 27 | path = os.path.join(model_dir, 'model_checkpoint') |
| 28 | if step is not None: |
| 29 | path += '-{:08d}'.format(step) |
| 30 | if os.path.exists(path): |
| 31 | print("Loading model from %s" % path) |
| 32 | checkpoint = torch.load(path, map_location=map_location) |
| 33 | |
| 34 | old_state_dict = item_dict["model"].state_dict() |
| 35 | for key in old_state_dict.keys(): |
| 36 | if key not in checkpoint['model']: |
| 37 | checkpoint['model'][key] = old_state_dict[key] |
| 38 | |
| 39 | for item_name in item_dict: |
| 40 | item_dict[item_name].load_state_dict(checkpoint[item_name]) |
| 41 | return checkpoint.get('step', 0) |
| 42 | return 0 |
| 43 | |
| 44 | |
| 45 | def load_and_map_checkpoint(model, model_dir, remap): |