File size: 3,148 Bytes
c319d57 |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import torch
from configuration_neuroclr import NeuroCLRConfig
from modeling_neuroclr import NeuroCLRForSequenceClassification
# -------- EDIT THESE PATHS + nhead if needed ----------
PRETRAIN_CKPT = ""
HEAD_CKPT = ""
OUT_DIR = "."
CFG = dict(
# encoder MUST match the pretrained export
TSlength=128,
nhead=2, # change if needed
nlayer=2, # we confirmed this from your pretraining ckpt
projector_out1=128,
projector_out2=64,
pooling="flatten",
normalize_input=True,
# classification
n_rois=200,
num_labels=2,
freeze_encoder=True, # encoder frozen by default
# ResNet1D head (your exact settings)
base_filters=256,
kernel_size=16,
stride=2,
groups=32,
n_block=48,
downsample_gap=6,
increasefilter_gap=12,
use_bn=True,
use_do=True,
)
# -----------------------------------------------------
def load_model_state_dict(path):
ckpt = torch.load(path, map_location="cpu")
if isinstance(ckpt, dict):
if "model_state_dict" in ckpt:
return ckpt["model_state_dict"]
if "state_dict" in ckpt:
return ckpt["state_dict"]
return ckpt
return ckpt
def remap_encoder(sd):
# pretraining ckpt keys: transformer_encoder.* and projector.*
new = {}
for k, v in sd.items():
k2 = k.replace("module.", "")
if k2.startswith("transformer_encoder.") or k2.startswith("projector."):
new["encoder." + k2] = v
return new
def remap_head(sd):
# head ckpt keys likely start with first_block_conv.*, basicblock_list.*, dense.* etc.
new = {}
for k, v in sd.items():
k2 = k.replace("module.", "")
head_prefixes = (
"first_block_conv.", "first_block_bn.", "first_block_relu.",
"basicblock_list.", "final_bn.", "final_relu.", "dense."
)
if k2.startswith(head_prefixes):
new["head." + k2] = v
# If your checkpoint already has head.* then keep it
elif k2.startswith("head."):
new[k2] = v
return new
def main():
config = NeuroCLRConfig(**CFG)
# Enables HF auto-classes loading from this folder
config.auto_map = {
"AutoConfig": "configuration_neuroclr.NeuroCLRConfig",
"AutoModelForSequenceClassification": "modeling_neuroclr.NeuroCLRForSequenceClassification",
}
model = NeuroCLRForSequenceClassification(config)
# 1) Load encoder weights from pretraining ckpt
enc_sd_raw = load_model_state_dict(PRETRAIN_CKPT)
enc_sd = remap_encoder(enc_sd_raw)
# 2) Load head weights from classification ckpt
head_sd_raw = load_model_state_dict(HEAD_CKPT)
head_sd = remap_head(head_sd_raw)
# 3) Merge and load
merged = {}
merged.update(enc_sd)
merged.update(head_sd)
missing, unexpected = model.load_state_dict(merged, strict=False)
print("Missing:", missing)
print("Unexpected:", unexpected)
# Save to HF folder
model.save_pretrained(OUT_DIR, safe_serialization=True)
print("Saved HF classification model to:", OUT_DIR)
if __name__ == "__main__":
main()
|