rpotter6298
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only."""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from v4.classes.accessory.se_block import SEBlock
|
|
|
|
|
|
class FusionBridge(nn.Module):
|
|
"""Project N input embeddings to a shared dim, fuse via element-wise product.
|
|
|
|
Pure embedding producer — no classification head. Attach a head stage in
|
|
the pipeline config to produce logits.
|
|
|
|
Parameters
|
|
----------
|
|
input_dims : ordered list of input embedding dims
|
|
fusion_dim : projection / output dimension
|
|
use_se : SE gate on the fused vector
|
|
se_reduction : SE reduction factor
|
|
se_pre_norm : LayerNorm before each projection; else Identity
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dims: list[int],
|
|
fusion_dim: int = 256,
|
|
use_se: bool = True,
|
|
se_reduction: int = 16,
|
|
se_pre_norm: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.out_dim = fusion_dim
|
|
self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
|
|
self.ln = nn.ModuleList(
|
|
[nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
|
for _ in input_dims]
|
|
)
|
|
self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
|
|
|
|
def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
|
|
assert len(embeddings) == len(self.W), (
|
|
f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
|
|
)
|
|
h = self.ln[0](self.W[0](embeddings[0]))
|
|
for i in range(1, len(embeddings)):
|
|
h = h * self.ln[i](self.W[i](embeddings[i]))
|
|
if self.se is not None:
|
|
h, _ = self.se(h)
|
|
return h
|
|
|
|
def set_phase(self, phase: str) -> None:
|
|
"""Freeze bridge during tower_warmup; trainable otherwise."""
|
|
enabled = phase not in ("tower_warmup", "cd_warmup")
|
|
for p in self.parameters():
|
|
p.requires_grad_(enabled)
|