24 lines
969 B
Python
24 lines
969 B
Python
import os
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
def load_state_dict(model, ckpt_path, device="cuda", strict=False, assign=True):
|
|
"""Load state dict from checkpoint file."""
|
|
if os.path.splitext(ckpt_path)[1] == ".safetensors":
|
|
state_dict = load_file(ckpt_path, device)
|
|
else:
|
|
state_dict = torch.load(ckpt_path, map_location="cpu")
|
|
|
|
missing, unexpected = model.load_state_dict(
|
|
state_dict, strict=strict, assign=assign
|
|
)
|
|
if len(missing) > 0 and len(unexpected) > 0:
|
|
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
|
print("\n" + "-" * 79 + "\n")
|
|
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
|
elif len(missing) > 0:
|
|
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
|
elif len(unexpected) > 0:
|
|
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
|
return model |