Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

belka_gate_fusions: inject GateFusionResidual variants to asym and dual encoder

Browse Source
This commit is contained in:
pikaliov
2026-05-18 11:24:58 +03:00
parent e2fd16e66c
commit b7542b27de
4 changed files with 167 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/models/asymmetric_encoder.py
View File

@@ -260,6 +260,9 @@ class TextFusionMLP(nn.Module):
# Main model: AsymmetricEncoder
# ---------------------------------------------------------------------------
# ResidualGateFusin experiment
from residual_fusions import ResidualGateType, GatedFusionResidual
class AsymmetricEncoder(nn.Module):
"""Dual encoder for CVGL with text fusion on both branches.
@@ -294,6 +297,7 @@ class AsymmetricEncoder(nn.Module):
def __init__(
self,
gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate,
dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
@@ -365,8 +369,12 @@ class AsymmetricEncoder(nn.Module):
)
# Separate gated fusion for query and gallery branches.
self.fusion_query = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
self.fusion_gallery = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
#! Experimental Gated fusion on query branch.
self.fusion_query = GatedFusionResidual(gate_type=gate_type,
init_gate=init_gate, baseline_mode=baseline_mode)
self.fusion_gallery = GatedFusionResidual(gate_type=gate_type,
init_gate=init_gate, baseline_mode=baseline_mode)
@staticmethod
def _freeze(module: nn.Module) -> None:
@@ -420,7 +428,7 @@ class AsymmetricEncoder(nn.Module):
l1_texts: list[str] | None,
l2_texts: list[str] | None,
l3_texts: list[str] | None,
fusion: GatedFusion,
fusion: GatedFusionResidual,
) -> torch.Tensor:
"""Fuse image features with optional text, respecting per-sample presence.
@@ -451,8 +459,7 @@ class AsymmetricEncoder(nn.Module):
# Per-sample fusion: text-present samples use full gated fusion,
# empty-caption samples pass through pure image features.
gate = torch.sigmoid(fusion.alpha)
fused_with_text = gate * img_feat + (1.0 - gate) * z_text
fused_with_text = fusion(img_feat, z_text)
out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat)
return F.normalize(out, dim=-1)

124
src/models/dual_encoder.py
View File

@@ -21,6 +21,9 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
# residual fusions exp
from residual_fusions import ResidualGateType, GatedFusionResidual
class ProjectionHead(nn.Module):
"""MLP projection head with L2 normalization."""
@@ -83,120 +86,6 @@ class GatedFusion(nn.Module):
return torch.sigmoid(self.alpha).item()
#! GATE-FUSIONS MODIFICATIONS ---------------------------------
#! in_dim = 1024
from enum import Enum
class ResidualGateType(Enum):
simple_residual_one_gate = 0,
cross_gate = 1,
gate_sum = 2,
alpha_res_cat = 3,
alpha_res_sum = 4
# TODO: add GatedFusionresidual class to gin
@gin.configurable
class GatedFusionresidual(nn.Module):
"""Learnable gated fusion of image and text embeddings.
V1 - Simple residual gating with 1 common gate:
V2 - Cross residual gating with 2 cross-gates
V3 - Gate + Simple Sum of feats x & y
V4 - Alpha-weighted residual concat (per sample)
V5 - Alpha-weighted residual sum (per sample)
"""
def __init__(self, init_gate: float = 0.7, in_dim = 1024, baseline_mode: bool = False) -> None:
super().__init__()
# alpha is in logit space: sigmoid(alpha) = init_gate
init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha = nn.Parameter(init_alpha)
# alphas for separated cases
init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha_img = nn.Parameter(init_alpha_img_cross_gate)
self.alpha_text = nn.Parameter(init_alpha_text_cross_gate)
# weight for sum and cat residual
self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim)
self.weight_net_for_sum = nn.Linear(in_dim, 1)
self.weight_net_for_cat = nn.Linear(in_dim * 2, 1)
self.baseline_mode = baseline_mode
def FuseSRGF(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
img_res = img_feat * gate + img_feat
text_res = text_feat * (1 - gate) + text_feat
fused_vec = img_res + text_res
return fused_vec
def FuseRCGF(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate_img = torch.sigmoid(self.alpha_img)
gate_text = torch.sigmoid(self.alpha_text)
z_img = img_feat + gate_text * img_feat
z_text = text_feat + gate_img * text_feat
fused_vec = z_img + z_text
return fused_vec
def FuseGSUM(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat
return fuzed_vec
def FuseARGFSum(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
residual = img_feat + text_feat
res_weight = torch.sigmoid(self.weight_net_for_sum(residual))
fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
return fuzed_vec
def FuseARGFCat(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
cat_vec = torch.cat([img_feat, text_feat], dim=-1)
residual = self.final_cat_residual_proj(cat_vec)
res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec))
fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
return fuzed_vec
def forward(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
if text_feat is None or self.baseline_mode:
return img_feat
gate = torch.sigmoid(self.alpha)
# TODO: switch forwards here
return gate * img_feat + (1.0 - gate) * text_feat
@gin.configurable
class DualEncoderCaptionTest(nn.Module):
"""GeoRSCLIP dual encoder with gated text fusion on query branch.
@@ -222,6 +111,7 @@ class DualEncoderCaptionTest(nn.Module):
baseline_mode: bool = False,
init_gate: float = 0.7,
device: str = "cuda",
gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate
) -> None:
super().__init__()
self.embed_dim = embed_dim
@@ -247,7 +137,11 @@ class DualEncoderCaptionTest(nn.Module):
self._apply_unfreeze(unfreeze_mode)
# Gated fusion on query branch.
self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
# self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
#! Experimental Gated fusion on query branch.
self.fusion = GatedFusionResidual(gate_type=gate_type,
init_gate=init_gate, baseline_mode=baseline_mode)
# Projection heads.
self.proj_query = ProjectionHead(

144
src/models/residual_fusions.py Normal file
View File

@@ -0,0 +1,144 @@
import torch
import torch.nn as nn
import gin
#! GATE-FUSIONS MODIFICATIONS ---------------------------------
#! in_dim = 1024
from enum import Enum
import math
class ResidualGateType(Enum):
simple_residual_one_gate = 0,
cross_gate = 1,
gate_sum = 2,
alpha_res_cat = 3,
alpha_res_sum = 4
# TODO: add GatedFusionresidual class to gin
def init_bias_for_sigmoid(linear: nn.Linear, value: float) -> None:
nn.init.zeros_(linear.weight)
nn.init.constant_(linear.bias, math.log(value / (1.0 - value)))
def init_residual_projs(linear: nn.Linear, scale: float) -> None:
nn.init.xavier_uniform_(linear.weight, gain=scale)
nn.init.zeros_(linear.bias)
RESIDUAL_GATES = {
ResidualGateType.alpha_res_sum,
ResidualGateType.alpha_res_cat
}
@gin.configurable
class GatedFusionResidual(nn.Module):
"""Learnable gated fusion of image and text embeddings.
V1 - Simple residual gating with 1 common gate:
V2 - Cross residual gating with 2 cross-gates
V3 - Gate + Simple Sum of feats x & y
V4 - Alpha-weighted residual sum (per sample)
V5 - Alpha-weighted residual concat (per sample)
"""
def __init__(self, gate_type: ResidualGateType,
init_gate: float = 0.7, in_dim = 1024,
init_res_weight: float = 0.1, residual_proj_scale: float = 0.1,
baseline_mode: bool = False,
) -> None:
super().__init__()
# alpha is in logit space: sigmoid(alpha) = init_gate
init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha = nn.Parameter(init_alpha)
# alphas for separated cases
if gate_type == ResidualGateType.cross_gate:
init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
self.alpha_img = nn.Parameter(init_alpha_img_cross_gate)
self.alpha_text = nn.Parameter(init_alpha_text_cross_gate)
# weight for sum and cat residual
if gate_type in RESIDUAL_GATES:
self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim)
self.weight_net_for_sum = nn.Linear(in_dim, 1)
self.weight_net_for_cat = nn.Linear(in_dim * 2, 1)
init_bias_for_sigmoid(self.weight_net_for_sum, value=init_res_weight)
init_bias_for_sigmoid(self.weight_net_for_cat, value=init_res_weight)
init_residual_projs(self.final_cat_residual_proj, scale=residual_proj_scale)
self.gate_type = gate_type
self.baseline_mode = baseline_mode
def FuseSRGF(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
img_res = img_feat * gate + img_feat
text_res = text_feat * (1 - gate) + text_feat
fused_vec = img_res + text_res
return fused_vec
def FuseRCGF(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate_img = torch.sigmoid(self.alpha_img)
gate_text = torch.sigmoid(self.alpha_text)
z_img = img_feat + gate_text * img_feat
z_text = text_feat + gate_img * text_feat
fused_vec = z_img + z_text
return fused_vec
def FuseGSUM(self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat
return fuzed_vec
def FuseARGFSum(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
residual = img_feat + text_feat
res_weight = torch.sigmoid(self.weight_net_for_sum(residual))
fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
return fuzed_vec
def FuseARGFCat(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
gate = torch.sigmoid(self.alpha)
cat_vec = torch.cat([img_feat, text_feat], dim=-1)
residual = self.final_cat_residual_proj(cat_vec)
res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec))
fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
return fuzed_vec
def forward(
self,
img_feat: torch.Tensor,
text_feat: torch.Tensor | None,
) -> torch.Tensor:
if text_feat is None or self.baseline_mode:
return img_feat
if self.gate_type == ResidualGateType.simple_residual_one_gate:
fused_vec = self.FuseSRGF(img_feat=img_feat, text_feat=text_feat)
if self.gate_type == ResidualGateType.cross_gate:
fused_vec = self.FuseRCGF(img_feat=img_feat, text_feat=text_feat)
if self.gate_type == ResidualGateType.gate_sum:
fused_vec = self.FuseGSUM(img_feat=img_feat, text_feat=text_feat)
if self.gate_type == ResidualGateType.alpha_res_sum:
fused_vec = self.FuseARGFSum(img_feat=img_feat, text_feat=text_feat)
if self.gate_type == ResidualGateType.alpha_res_cat:
fused_vec = self.FuseARGFCat(img_feat=img_feat, text_feat=text_feat)
return fused_vec

3
src/training/train_gtauav.py
View File

@@ -99,7 +99,8 @@ class TrainConfigGTAUAV:
# Training.
resume_from: str | None = None # path to checkpoint for resuming
output_dir: str = "out/gtauav/with_text"
# output_dir: str = "out/gtauav/with_text"
output_dir: str = "out/gtauav/with_text_exp_gate_SRGF"
epochs: int = 10
batch_size: int = 8
num_workers: int = 4