feat(phase4c): weight loading + pos_embed bicubic interpolation
HieraBackbone.load_checkpoint() loads safetensors, strips encoder.model. prefix, bicubic-interpolates pos_embed from 512x512 to 128x128 tokens. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
853efbe2cf
commit
0d6d41cb45
@ -9,10 +9,16 @@ Reference: https://arxiv.org/abs/2306.00989
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from functools import reduce
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from safetensors import safe_open
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
# ── hiera_base_plus_224 constants ──────────────────────────────────────
|
||||
EMBED_DIM = 112
|
||||
@ -29,6 +35,8 @@ IN_CHANS = 4 # RGB + alpha hint
|
||||
MLP_RATIO = 4.0
|
||||
DIM_MUL = 2.0
|
||||
HEAD_MUL = 2.0
|
||||
ENCODER_KEY_PREFIX = "encoder.model."
|
||||
TRAIN_IMG_SIZE = 2048 # checkpoint was trained at this resolution
|
||||
|
||||
|
||||
# ── Helpers ────────────────────────────────────────────────────────────
|
||||
@ -161,6 +169,36 @@ def reroll(
|
||||
return undo_windowing(x, size, cur_mu_shape)
|
||||
|
||||
|
||||
def _interpolate_pos_embed(
|
||||
ckpt_embed: mx.array,
|
||||
target_tokens: int,
|
||||
) -> mx.array:
|
||||
"""Bicubic interpolation of pos_embed from checkpoint to model resolution.
|
||||
|
||||
Args:
|
||||
ckpt_embed: (1, N_ckpt, C) from checkpoint
|
||||
target_tokens: target token count N_model = H_model * W_model
|
||||
|
||||
Returns:
|
||||
(1, N_model, C)
|
||||
"""
|
||||
ckpt_n = ckpt_embed.shape[1]
|
||||
if ckpt_n == target_tokens:
|
||||
return ckpt_embed
|
||||
|
||||
embed_dim = ckpt_embed.shape[2]
|
||||
ckpt_side = int(math.sqrt(ckpt_n))
|
||||
model_side = int(math.sqrt(target_tokens))
|
||||
|
||||
# (1, N, C) -> (1, H, W, C) NHWC for MLX upsample
|
||||
embed = ckpt_embed.reshape(1, ckpt_side, ckpt_side, embed_dim)
|
||||
scale = model_side / ckpt_side
|
||||
resizer = nn.Upsample(scale_factor=(scale, scale), mode="cubic", align_corners=False)
|
||||
embed = resizer(embed)
|
||||
# Back to (1, N, C)
|
||||
return embed.reshape(1, target_tokens, embed_dim)
|
||||
|
||||
|
||||
# ── Modules ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@ -383,6 +421,34 @@ class HieraBackbone(nn.Module):
|
||||
self.blocks.append(block)
|
||||
embed_dim = dim_out
|
||||
|
||||
def load_checkpoint(self, path: str | Path) -> None:
|
||||
"""Load weights from converted safetensors checkpoint.
|
||||
|
||||
Strips ``encoder.model.`` prefix and bicubic-interpolates pos_embed
|
||||
from training resolution to model resolution.
|
||||
"""
|
||||
target_tokens = _prod(self.tokens_spatial_shape)
|
||||
weight_pairs: list[tuple[str, mx.array]] = []
|
||||
|
||||
with safe_open(str(path), framework="numpy") as f:
|
||||
for full_key in f.keys(): # noqa: SIM118 — safe_open isn't iterable
|
||||
if not full_key.startswith(ENCODER_KEY_PREFIX):
|
||||
continue
|
||||
mlx_key = full_key[len(ENCODER_KEY_PREFIX) :]
|
||||
tensor = mx.array(f.get_tensor(full_key))
|
||||
|
||||
if mlx_key == "pos_embed":
|
||||
tensor = _interpolate_pos_embed(tensor, target_tokens)
|
||||
# materialize interpolated embedding
|
||||
mx.eval(tensor) # noqa: S307 — mx.eval, not Python eval
|
||||
|
||||
weight_pairs.append((mlx_key, tensor))
|
||||
|
||||
self.load_weights(weight_pairs)
|
||||
self.eval()
|
||||
# materialize all parameters
|
||||
mx.eval(self.parameters()) # noqa: S307 — mx.eval, not Python eval
|
||||
|
||||
def __call__(self, x: mx.array) -> list[mx.array]:
|
||||
"""Forward pass.
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user