(model, model_dir, remap)
| 43 | |
| 44 | |
| 45 | def load_and_map_checkpoint(model, model_dir, remap): |
| 46 | path = os.path.join(model_dir, 'model_checkpoint') |
| 47 | print("Loading parameters %s from %s" % (remap.keys(), model_dir)) |
| 48 | checkpoint = torch.load(path) |
| 49 | new_state_dict = model.state_dict() |
| 50 | for name, value in remap.items(): |
| 51 | # TODO: smarter mapping. |
| 52 | new_state_dict[name] = checkpoint['model'][value] |
| 53 | model.load_state_dict(new_state_dict) |
| 54 | |
| 55 | |
| 56 | def save_checkpoint(items, step, model_dir, ignore=[], |