| 63 | |
| 64 | |
| 65 | def load_model_weights(model, ckpt_path, strict=True, verbose=True): |
| 66 | ckpt = torch.load(ckpt_path, map_location="cpu") |
| 67 | if verbose: |
| 68 | print("-" * 40) |
| 69 | print(f"{'Loading weights':<16}: {ckpt_path}") |
| 70 | print(f"{'Model':<16}: {model.__class__.__name__}") |
| 71 | if "global_step" in ckpt: |
| 72 | print(f"{f'Global Step':<16}: {ckpt['global_step']:,}") |
| 73 | print(f"{'Strict':<16}: {'True' if strict else 'False'}") |
| 74 | print("-" * 40) |
| 75 | sd = ckpt["state_dict"] if 'state_dict' in ckpt else ckpt |
| 76 | missing, unexpected = model.load_state_dict(sd, strict=strict) |
| 77 | if len(missing) > 0: |
| 78 | warnings.warn(f"Load model weights - missing keys: {len(missing)}") |
| 79 | if verbose: |
| 80 | print(missing) |
| 81 | if len(unexpected) > 0: |
| 82 | warnings.warn(f"Load model weights - unexpected keys: {len(unexpected)}") |
| 83 | if verbose: |
| 84 | print(unexpected) |
| 85 | return model |
| 86 | |
| 87 | def load_lora_weights(model, ckpt_path, strict=False, verbose=False): |
| 88 | ckpt = torch.load(ckpt_path, map_location="cpu") |