feat: add mamba and dynamic chunking related code and test code
This commit is contained in:
gameloader
339
models/DC_hnet.py
Normal file
339
models/DC_hnet.py
Normal file
@ -0,0 +1,339 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Literal, List
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# 来自你的代码库（可直接使用）
|
||||
from hnet.modules.dc import RoutingModule, ChunkLayer
|
||||
from hnet.modules.isotropic import Isotropic
|
||||
from hnet.models.config_hnet import HNetConfig, SSMConfig, AttnConfig
|
||||
|
||||
# -------------------- 辅助 --------------------
|
||||
def create_isotropic_encoder(d_model, arch="m", height=4, device=None, dtype=None):
|
||||
"""创建简化的Isotropic编码器"""
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
|
||||
# 创建HNetConfig，确保list字段有足够的元素
|
||||
config = HNetConfig(
|
||||
arch_layout=[f"{arch}{height}"],
|
||||
d_model=[d_model],
|
||||
d_intermediate=[d_model * 2],
|
||||
ssm_cfg=SSMConfig(
|
||||
d_conv=4,
|
||||
expand=2,
|
||||
d_state=128,
|
||||
chunk_size=256
|
||||
),
|
||||
attn_cfg=AttnConfig(
|
||||
num_heads=[8], # 确保有至少一个元素
|
||||
rotary_emb_dim=[0], # 确保有至少一个元素
|
||||
window_size=[-1] # 确保有至少一个元素
|
||||
)
|
||||
)
|
||||
|
||||
return Isotropic(
|
||||
config=config,
|
||||
pos_idx=0,
|
||||
stage_idx=0,
|
||||
**factory_kwargs
|
||||
)
|
||||
|
||||
def ratio_loss(boundary_mask: torch.Tensor, boundary_prob: torch.Tensor, target_N: int) -> torch.Tensor:
|
||||
F_act = boundary_mask.float().mean(dim=1) # (B,)
|
||||
G_prob = boundary_prob[..., 1].mean(dim=1) # (B,)
|
||||
N = float(target_N)
|
||||
loss = N / (N - 1.0) * (((N - 1.0) * F_act) + (1.0 - F_act) * (1.0 - G_prob))
|
||||
return loss.mean()
|
||||
|
||||
def masked_mean(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
||||
mask_f = mask.float().unsqueeze(-1) # (B, L, 1)
|
||||
s = (x * mask_f).sum(dim=1) # (B, D)
|
||||
denom = mask_f.sum(dim=1).clamp_min(1.0)
|
||||
return s / denom
|
||||
|
||||
# -------------------- 多层Encoder（金字塔）：每层Mamba2 + 路由下采样，只有最终有主网络 --------------------
|
||||
class PyramidEncoders_NoDechunk(nn.Module):
|
||||
"""
|
||||
层级结构（仅编码器逐层压缩；主网络只在最终一层）：
|
||||
输入 x0: (B, L0, 1)
|
||||
- 线性升维 -> D0
|
||||
For s = 0..S-1:
|
||||
Es(Mamba2, D_s) -> h_s (B, L_s, D_s)
|
||||
路由 + 下采样 -> x_{s+1} (B, L_{s+1}, D_s), mask_{s+1}
|
||||
维度扩展 D_s -> D_{s+1}（拼接共享向量）
|
||||
最终 x_S: (B, L_S, D_S) 送入单一主网络 M (Transformer/Mamba)
|
||||
跨尺度融合（不去分块）：融合 E^0 的 pooled_enc0 与 主网络 pooled_main
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
d_models: List[int], # [D0, D1, ..., D_S] 单调非降
|
||||
encoder_cfg_per_stage: List[dict], # S个编码器配置（必须 arch='m'/'M'）
|
||||
main_cfg: dict, # 单一主网络配置（在最压缩序列上工作）
|
||||
fusion_dropout: float = 0.1,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
super().__init__()
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
|
||||
assert len(d_models) >= 1
|
||||
S = len(d_models) - 1
|
||||
assert S == len(encoder_cfg_per_stage), "stage数等于encoder配置数"
|
||||
for i in range(S):
|
||||
assert d_models[i+1] >= d_models[i], "需满足 D_s <= D_{s+1}（宽度单调增加）"
|
||||
assert encoder_cfg_per_stage[i].get("arch", "m") in ("m", "M"), "Encoder必须为Mamba2"
|
||||
|
||||
self.S = S
|
||||
self.d_models = d_models
|
||||
|
||||
# 输入升维到 D0
|
||||
self.input_proj = nn.Linear(1, d_models[0], **factory_kwargs)
|
||||
|
||||
# 每层编码器 + 路由 + 下采样 + 扩宽参数
|
||||
self.encoders = nn.ModuleList()
|
||||
self.routers = nn.ModuleList()
|
||||
self.chunks = nn.ModuleList()
|
||||
self.pad_vectors = nn.ParameterList()
|
||||
for s in range(S):
|
||||
self.encoders.append(
|
||||
create_isotropic_encoder(
|
||||
d_model=d_models[s],
|
||||
**{k: v for k, v in encoder_cfg_per_stage[s].items() if k != "d_model"},
|
||||
**factory_kwargs
|
||||
)
|
||||
)
|
||||
self.routers.append(RoutingModule(d_models[s], **factory_kwargs))
|
||||
self.chunks.append(ChunkLayer())
|
||||
delta = d_models[s+1] - d_models[s]
|
||||
self.pad_vectors.append(nn.Parameter(torch.zeros(delta, **factory_kwargs)) if delta > 0 else nn.Parameter(torch.empty(0, **factory_kwargs)))
|
||||
|
||||
# 最终唯一的主网络：在 D_S & L_S 上运行
|
||||
self.main_network = create_isotropic_encoder(
|
||||
d_model=d_models[-1],
|
||||
**{k: v for k, v in main_cfg.items() if k != "d_model"},
|
||||
**factory_kwargs
|
||||
)
|
||||
|
||||
# 跨尺度融合：将 pooled_enc0(D0) 投到 D_S 并与 pooled_main(D_S) 融合 -> D_S
|
||||
self.proj_enc0_to_DS = nn.Linear(d_models[0], d_models[-1], **factory_kwargs)
|
||||
self.fusion_head = nn.Sequential(
|
||||
nn.Linear(d_models[-1] + d_models[-1], d_models[-1], **factory_kwargs),
|
||||
nn.GELU(),
|
||||
nn.Dropout(fusion_dropout),
|
||||
nn.Linear(d_models[-1], d_models[-1], **factory_kwargs),
|
||||
)
|
||||
|
||||
def _expand_width(self, x: torch.Tensor, pad_vec: nn.Parameter) -> torch.Tensor:
|
||||
if pad_vec.numel() == 0:
|
||||
return x
|
||||
early = x.shape[:-1]
|
||||
return torch.cat([x, pad_vec.expand(*early, -1)], dim=-1)
|
||||
|
||||
def forward(self, x_scalar: torch.Tensor, mask: Optional[torch.Tensor] = None, return_seq: bool = False):
|
||||
"""
|
||||
x_scalar: (B, L) 或 (B, L, 1)
|
||||
mask: (B, L) bool
|
||||
返回:
|
||||
fused_vec: (B, D_S)
|
||||
debug: 可选
|
||||
aux: 包含各层路由信息（供ratio loss）
|
||||
"""
|
||||
if x_scalar.dim() == 2:
|
||||
x_scalar = x_scalar.unsqueeze(-1) # (B, L, 1)
|
||||
B, L, _ = x_scalar.shape
|
||||
device = x_scalar.device
|
||||
if mask is None:
|
||||
mask = torch.ones(B, L, dtype=torch.bool, device=device)
|
||||
|
||||
# 初始升维到 D0
|
||||
x = self.input_proj(x_scalar) # (B, L0, D0)
|
||||
cur_mask = mask
|
||||
|
||||
pooled_enc0 = None
|
||||
aux_per_stage = []
|
||||
seq_debug = [] if return_seq else None
|
||||
|
||||
# 逐层：Encoder(Mamba2)->Routing->Chunk->Expand D
|
||||
for s in range(self.S):
|
||||
d_in = self.d_models[s]
|
||||
# 细粒度编码（未压缩序列）
|
||||
h_enc = self.encoders[s](x, mask=cur_mask) # (B, L_s, D_s)
|
||||
|
||||
if s == 0:
|
||||
pooled_enc0 = masked_mean(h_enc, cur_mask) # (B, D0)
|
||||
|
||||
# 路由 + 下采样（得到更短序列）
|
||||
bpred = self.routers[s](h_enc, mask=cur_mask)
|
||||
x_next, _, _, mask_next = self.chunks[s](h_enc, bpred.boundary_mask, mask=cur_mask) # (B, L_{s+1}, D_s)
|
||||
|
||||
# 扩展宽度 D_s -> D_{s+1}
|
||||
x_next = self._expand_width(x_next, self.pad_vectors[s]) # (B, L_{s+1}, D_{s+1})
|
||||
|
||||
# 推进到下一层
|
||||
x, cur_mask = x_next, mask_next
|
||||
|
||||
aux_per_stage.append({
|
||||
"boundary_mask": bpred.boundary_mask,
|
||||
"boundary_prob": bpred.boundary_prob,
|
||||
"selected_probs": bpred.selected_probs,
|
||||
})
|
||||
if return_seq:
|
||||
seq_debug.append({"stage": s, "seq": x, "mask": cur_mask})
|
||||
|
||||
# 现在 x: (B, L_S, D_S), cur_mask: (B, L_S)
|
||||
# 最终单一主网络在最压缩序列上
|
||||
h_main = self.main_network(x, mask=cur_mask) # (B, L_S, D_S)
|
||||
|
||||
# 主网络池化
|
||||
if cur_mask is None:
|
||||
pooled_main = h_main.mean(dim=1) # (B, D_S)
|
||||
else:
|
||||
pooled_main = (h_main * cur_mask.float().unsqueeze(-1)).sum(dim=1) / \
|
||||
cur_mask.float().sum(dim=1, keepdim=True).clamp_min(1.0)
|
||||
|
||||
# 跨尺度融合：E^0 全局池化 与 主网络池化
|
||||
pooled_enc0_in_DS = self.proj_enc0_to_DS(pooled_enc0) # (B, D_S)
|
||||
fused = torch.cat([pooled_enc0_in_DS, pooled_main], dim=-1) # (B, 2*D_S)
|
||||
fused = self.fusion_head(fused) # (B, D_S)
|
||||
|
||||
aux = {"per_stage": aux_per_stage}
|
||||
if return_seq:
|
||||
return fused, {"stages": seq_debug, "main_seq": h_main, "main_mask": cur_mask}, aux
|
||||
else:
|
||||
return fused, None, aux
|
||||
|
||||
# -------------------- 顶层：多通道融合 + 分类头（仅一个主网络） --------------------
|
||||
@dataclass
|
||||
class HierEncodersSingleMainConfig:
|
||||
num_channels: int
|
||||
d_models: List[int] # [D0, D1, ..., D_S] 单调非降
|
||||
num_classes: int
|
||||
encoder_cfg_per_stage: List[dict] # S个编码器配置（均为Mamba2, height≈4）
|
||||
main_cfg: dict # 单一主网络配置（Transformer或Mamba2），d_model自动用D_S
|
||||
target_compression_N_per_stage: List[int]
|
||||
share_channel: bool = True
|
||||
fusion_across_channels: Literal["mean", "concat"] = "mean"
|
||||
dropout: float = 0.1
|
||||
|
||||
class HierEncodersSingleMainClassifier(nn.Module):
|
||||
def __init__(self, cfg: HierEncodersSingleMainConfig, dtype=None, device=None):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
|
||||
S = len(cfg.d_models) - 1
|
||||
assert S == len(cfg.encoder_cfg_per_stage) == len(cfg.target_compression_N_per_stage), "stage数不一致"
|
||||
|
||||
if cfg.share_channel:
|
||||
self.channel_encoder = PyramidEncoders_NoDechunk(
|
||||
d_models=cfg.d_models,
|
||||
encoder_cfg_per_stage=cfg.encoder_cfg_per_stage,
|
||||
main_cfg=cfg.main_cfg,
|
||||
**factory_kwargs,
|
||||
)
|
||||
else:
|
||||
self.channel_encoder = nn.ModuleList([
|
||||
PyramidEncoders_NoDechunk(
|
||||
d_models=cfg.d_models,
|
||||
encoder_cfg_per_stage=cfg.encoder_cfg_per_stage,
|
||||
main_cfg=cfg.main_cfg,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(cfg.num_channels)
|
||||
])
|
||||
|
||||
fusion_dim = (cfg.num_channels * cfg.d_models[-1]) if cfg.fusion_across_channels == "concat" \
|
||||
else cfg.d_models[-1]
|
||||
self.dropout = nn.Dropout(cfg.dropout)
|
||||
self.head = nn.Linear(fusion_dim, cfg.num_classes, **factory_kwargs)
|
||||
|
||||
def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, return_seq: bool = False):
|
||||
"""
|
||||
x: (B, L, N) 多通道输入
|
||||
mask: (B, L) 时序mask
|
||||
"""
|
||||
B, L, N = x.shape
|
||||
assert N == self.cfg.num_channels
|
||||
|
||||
channel_vecs: List[torch.Tensor] = []
|
||||
ratio_losses = []
|
||||
seq_dbg_all = [] if return_seq else None
|
||||
|
||||
for c in range(N):
|
||||
x_c = x[..., c] # (B, L)
|
||||
if self.cfg.share_channel:
|
||||
vec, seq_dbg, aux = self.channel_encoder(x_c, mask=mask, return_seq=return_seq)
|
||||
else:
|
||||
vec, seq_dbg, aux = self.channel_encoder[c](x_c, mask=mask, return_seq=return_seq)
|
||||
|
||||
# ratio loss 累加（每个encoder stage一项）
|
||||
total_rl = 0.0
|
||||
for s, aux_s in enumerate(aux["per_stage"]):
|
||||
rl = ratio_loss(aux_s["boundary_mask"], aux_s["boundary_prob"], self.cfg.target_compression_N_per_stage[s])
|
||||
total_rl = total_rl + rl
|
||||
ratio_losses.append(total_rl)
|
||||
|
||||
channel_vecs.append(vec)
|
||||
if return_seq:
|
||||
seq_dbg_all.append(seq_dbg)
|
||||
|
||||
if self.cfg.fusion_across_channels == "concat":
|
||||
fused = torch.cat(channel_vecs, dim=-1) # (B, N*D_S)
|
||||
else:
|
||||
fused = torch.stack(channel_vecs, dim=1).mean(dim=1) # (B, D_S)
|
||||
|
||||
fused = self.dropout(fused)
|
||||
logits = self.head(fused)
|
||||
|
||||
aux_all = {"ratio_loss": torch.stack(ratio_losses).mean()}
|
||||
if return_seq:
|
||||
return logits, seq_dbg_all, aux_all
|
||||
else:
|
||||
return logits, None, aux_all
|
||||
|
||||
# -------------------- 使用示例 --------------------
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
符合要求：
|
||||
- 多层仅增加编码器数量（每层Mamba2 + 动态分块），主网络只有最终一个
|
||||
- 序列长度逐层缩短（由DC决定），通道维度 d_model 单调增大（SpaceByte式共享向量拼接）
|
||||
- 不使用去分块（dechunk）；跨尺度融合用 E^0 的全局池化 + 最终主网络池化
|
||||
"""
|
||||
B, L, N = 8, 1024, 6
|
||||
num_classes = 7
|
||||
d_models = [128, 256, 512] # D0 <= D1 <= D2
|
||||
|
||||
encoder_cfg_per_stage = [
|
||||
dict(arch="m", height=4, ssm_cfg=dict(), attn_cfg=dict()), # stage 0 encoder (Mamba2)
|
||||
dict(arch="m", height=4, ssm_cfg=dict(), attn_cfg=dict()), # stage 1 encoder (Mamba2)
|
||||
]
|
||||
main_cfg = dict(
|
||||
arch="T", height=12, ssm_cfg=dict(), attn_cfg=dict(num_heads=8) # 最终主网络（较重）
|
||||
)
|
||||
target_compression_N_per_stage = [4, 4]
|
||||
|
||||
cfg = HierEncodersSingleMainConfig(
|
||||
num_channels=N,
|
||||
d_models=d_models,
|
||||
num_classes=num_classes,
|
||||
encoder_cfg_per_stage=encoder_cfg_per_stage,
|
||||
main_cfg=main_cfg,
|
||||
target_compression_N_per_stage=target_compression_N_per_stage,
|
||||
share_channel=True,
|
||||
fusion_across_channels="mean",
|
||||
dropout=0.1,
|
||||
)
|
||||
|
||||
model = HierEncodersSingleMainClassifier(cfg).cuda().train()
|
||||
x = torch.randn(B, L, N, device="cuda")
|
||||
mask = torch.ones(B, L, dtype=torch.bool, device="cuda")
|
||||
|
||||
logits, _, aux = model(x, mask=mask, return_seq=False)
|
||||
y = torch.randint(0, num_classes, (B,), device="cuda")
|
||||
cls_loss = F.cross_entropy(logits, y)
|
||||
ratio_reg = 0.03 * aux["ratio_loss"]
|
||||
loss = cls_loss + ratio_reg
|
||||
loss.backward()
|
||||
print("logits:", logits.shape, "loss:", float(loss))
|
||||
Reference in New Issue
Block a user