File size: 5,189 Bytes
e379257 | 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 | """Student = Qwen3.5 backbone + LoRA + Clef-style JointSchemaHead (option logits -> softmax per question)
(+ optional BoxHead for Jev-style `box` questions, see glint/box.py)."""
from __future__ import annotations
import json, os
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
from peft import LoraConfig, PeftModel, get_peft_model
from .joint_schema_model import JointSchemaHead
from .box import BoxHead
from .glint_model import GlintModel
# LoRAは言語モデル側のみ（ビジョンエンコーダは凍結）。Gated DeltaNet層の射影も含める。
LORA_TARGET = (r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj"
r"|in_proj_qkv|in_proj_z|out_proj)")
def default_device():
return "cuda" if torch.cuda.is_available() else "cpu"
def parse_dtype(name: str):
return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[name]
def load_processor(base: str, max_pixels: int, min_pixels: int = 65536):
proc = AutoProcessor.from_pretrained(base)
ip = proc.image_processor
ip.size = {"shortest_edge": min_pixels, "longest_edge": max_pixels}
if hasattr(ip, "max_pixels"):
ip.max_pixels = max_pixels
if hasattr(ip, "min_pixels"):
ip.min_pixels = min_pixels
return proc
def build_student(base="Qwen/Qwen3.5-2B", max_pixels=430080, device=None, dtype="bf16",
lora_r=32, lora_alpha=64, lora_dropout=0.05, head_cfg=None, grad_ckpt=False, box_head_cfg=None):
"""box_head_cfg: None = no box head (previous behaviour); a dict (may be empty) = BoxHead kwargs."""
device = device or default_device()
td = parse_dtype(dtype)
proc = load_processor(base, max_pixels)
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(base, dtype=td).to(device)
backbone.config.use_cache = False
if grad_ckpt:
backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
backbone.enable_input_require_grads()
backbone = get_peft_model(backbone, LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
target_modules=LORA_TARGET))
hidden = backbone.config.text_config.hidden_size
head_cfg = dict(head_cfg or {"width": 512, "routing_layers": 2, "layers": 2, "heads": 8, "feedforward": 2048})
head_cfg["hidden_size"] = hidden
head = JointSchemaHead(**head_cfg, dropout=0.1).to(device, td)
box_head = None
if box_head_cfg is not None: # built after the joint head: its init RNG is unchanged
box_head_cfg = {**box_head_cfg, "hidden_size": hidden}
box_head = BoxHead(**box_head_cfg, dropout=0.1).to(device, td)
model = GlintModel(backbone, head, box_head)
meta = {"base_model": base, "max_pixels": max_pixels, "dtype": dtype, "head_config": head_cfg,
"lora": {"r": lora_r, "alpha": lora_alpha, "dropout": lora_dropout, "target_modules": LORA_TARGET}}
if box_head_cfg is not None:
meta["box_head_config"] = box_head_cfg
return model, proc, meta
def save_checkpoint(model: GlintModel, meta: dict, out_dir: str | os.PathLike, extra: dict | None = None):
out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
model.language_model.save_pretrained(out / "adapter") # LoRAのみ
save_file({k: v.detach().cpu().contiguous() for k, v in model.head.state_dict().items()}, out / "head.safetensors")
if getattr(model, "box_head", None) is not None:
save_file({k: v.detach().cpu().contiguous() for k, v in model.box_head.state_dict().items()}, out / "box_head.safetensors")
(out / "glint_config.json").write_text(json.dumps({**meta, **(extra or {})}, indent=2, ensure_ascii=False), encoding="utf-8")
def load_checkpoint(ckpt_dir: str | os.PathLike, device=None, max_pixels: int | None = None, merge=True):
ckpt = Path(ckpt_dir); device = device or default_device()
meta = json.loads((ckpt / "glint_config.json").read_text(encoding="utf-8"))
td = parse_dtype(meta["dtype"])
proc = load_processor(meta["base_model"], max_pixels or meta["max_pixels"])
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(meta["base_model"], dtype=td).to(device)
backbone.config.use_cache = False
backbone = PeftModel.from_pretrained(backbone, ckpt / "adapter")
if merge:
backbone = backbone.merge_and_unload() # 推論時はLoRAをマージして高速化
head = JointSchemaHead(**meta["head_config"]).to(device, td)
head.load_state_dict(load_file(ckpt / "head.safetensors"), strict=True)
box_head = None
if meta.get("box_head_config") is not None: # older checkpoints have no box head
box_head = BoxHead(**meta["box_head_config"]).to(device, td)
box_head.load_state_dict(load_file(ckpt / "box_head.safetensors"), strict=True)
return GlintModel(backbone, head, box_head).eval(), proc, meta
|