(model: jt.Module, ckpt_dir, file_weight_map)
| 27 | # check_state_dict(model, ckpt_dir, file_weight_map) |
| 28 | |
| 29 | def load_from_map(model: jt.Module, ckpt_dir, file_weight_map): |
| 30 | |
| 31 | for filename, names in tqdm(file_weight_map.items()): |
| 32 | cur_state_dict = torch.load(os.path.join(ckpt_dir, filename)) |
| 33 | for key, value in cur_state_dict.items(): |
| 34 | var = jt.Var(value.numpy()) |
| 35 | if value.requires_grad: |
| 36 | var.start_grad() |
| 37 | else: |
| 38 | var.stop_grad() |
| 39 | cur_state_dict[key] = var |
| 40 | |
| 41 | model.load_state_dict(cur_state_dict) |
| 42 | |
| 43 | # gc to reduce memory usage |
| 44 | del cur_state_dict |
| 45 | jt.sync_all() |
| 46 | jt.gc() |
| 47 | |
| 48 | def check_state_dict(model: jt.Module, ckpt_dir, file_weight_map): |
| 49 | for filename, names in file_weight_map.items(): |
no outgoing calls
no test coverage detected