File size: 2,282 Bytes
2f6b87a 92a70a2 2f6b87a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Load a released checkpoint into the model defined by the code release's methods.py.
Place this file in the root directory of the code release, next to methods.py.
"""
import re
import torch
from huggingface_hub import hf_hub_download
import methods
REPO_ID = "headless-start/parameter-efficient-dfer"
CLASS_NAMES = ["angry", "disgust", "fear", "happy", "sad", "surprise"]
def _rename(name):
# The training runs used older parameter names than the code release.
name = re.sub(r"\.ssf_gamma$", ".gamma", name)
name = re.sub(r"\.ssf_beta$", ".beta", name)
for old, new in ((".lora_q_A.", ".q_a."), (".lora_q_B.", ".q_b."),
(".lora_v_A.", ".v_a."), (".lora_v_B.", ".v_b."),
(".adaptmlp.", ".adapter.")):
name = name.replace(old, new)
return name
def _state(path):
return torch.load(path, map_location="cpu", weights_only=False)["model"]
def load_stage1():
model = methods.build_model(pretrained=False)
model.load_state_dict(_state(hf_hub_download(REPO_ID, "stage1/best.ckpt")))
return model.eval()
def load_stage2(method, fold):
"""method: full_ft, linear_probe, ssf, lora or adaptformer; fold: 0-9."""
file = f"stage2/{method}/fold_{fold:02d}/best.ckpt"
state = {_rename(k): v for k, v in _state(hf_hub_download(REPO_ID, file)).items()}
model = methods.build_model(pretrained=False)
if method != "full_ft":
# Adaptation weights sit on top of the frozen Stage-1 encoder.
backbone = _state(hf_hub_download(REPO_ID, "stage1/best.ckpt"))
model.load_state_dict({k: v for k, v in backbone.items()
if not k.startswith("head.")}, strict=False)
methods.configure(model, method, rank=4, reduction=12)
_, unexpected = model.load_state_dict(state, strict=method == "full_ft")
if unexpected:
raise ValueError(f"unexpected keys in {file}: {unexpected[:3]}")
return model.eval()
if __name__ == "__main__":
import sys
from PIL import Image
from data import build_transform
model = load_stage2("lora", 0)
image = build_transform(False, kmufed=True)(Image.open(sys.argv[1]).convert("RGB"))
with torch.no_grad():
print(CLASS_NAMES[model(image[None]).argmax(1).item()])
|