rpotter6298
97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from v3.classes.SE_attention import SEBlock, SEGateLogger
|
|
|
|
|
|
class Bridge(nn.Module):
|
|
def __init__(
|
|
self,
|
|
img_dim,
|
|
meta_dim,
|
|
num_classes,
|
|
fusion_dim=256,
|
|
mode="fused",
|
|
dropout: float = 0.5,
|
|
use_se: bool = True,
|
|
se_reduction: int = 16,
|
|
se_pre_norm: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.mode = mode
|
|
self.use_se = use_se
|
|
|
|
# project towers to equal width
|
|
self.W_img = nn.Linear(img_dim, fusion_dim)
|
|
self.W_md = nn.Linear(meta_dim, fusion_dim)
|
|
|
|
# optional: layernorm before SE
|
|
self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
|
self.ln_md = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
|
|
|
# SE gate on the fused vector
|
|
self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
|
|
self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim)
|
|
|
|
# heads
|
|
self.classifier_fused = nn.Sequential(
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(fusion_dim, num_classes),
|
|
)
|
|
self.classifier_img = nn.Linear(img_dim, num_classes)
|
|
self.classifier_cd = nn.Linear(meta_dim, num_classes)
|
|
|
|
def reset_se_stats(self):
|
|
"""Call at epoch start."""
|
|
if getattr(self, "se_log", None):
|
|
self.se_log.reset()
|
|
|
|
def get_se_stats(self, reset: bool = True):
|
|
"""Call after eval. Returns dict or None."""
|
|
if getattr(self, "se_log", None) and self.se_log.enabled:
|
|
return self.se_log.get(reset=reset)
|
|
return None
|
|
|
|
def _compute_fused(self, img_feats, md_feats):
|
|
"""Return z_fused embedding (before classifier_fused). Used by encode() and forward()."""
|
|
hi = self.ln_img(self.W_img(img_feats))
|
|
hm = self.ln_md(self.W_md(md_feats))
|
|
fused = hi * hm
|
|
if self.se is not None:
|
|
fused, gates = self.se(fused)
|
|
if self.se_log.enabled:
|
|
self.se_log.accumulate(gates)
|
|
return fused
|
|
|
|
def encode(self, img_feats, md_feats) -> torch.Tensor:
|
|
"""Return z_fused embedding without applying the classifier head."""
|
|
assert self.mode == "fused", "encode() only valid in fused mode"
|
|
return self._compute_fused(img_feats, md_feats)
|
|
|
|
def forward(self, img_feats, md_feats):
|
|
out_img = None if self.mode == "clinical_only" else self.classifier_img(img_feats)
|
|
out_md = None if self.mode == "image_only" else self.classifier_cd(md_feats)
|
|
|
|
if self.mode == "fused":
|
|
fused = self._compute_fused(img_feats, md_feats)
|
|
out_f = self.classifier_fused(fused)
|
|
return out_f, out_img, out_md
|
|
# if ablation modes:
|
|
if self.mode == "image_only":
|
|
return out_img, out_img, None
|
|
if self.mode == "clinical_only":
|
|
return out_md, None, out_md
|
|
|
|
|
|
class VoteBridge(nn.Module):
|
|
def __init__(self, num_classes):
|
|
super().__init__()
|
|
self.vote_combiner = nn.Linear(num_classes * 2, num_classes) # two sets of logits
|
|
|
|
def forward(self, out_img, out_md):
|
|
votes = torch.cat([out_img, out_md], dim=1)
|
|
return self.vote_combiner(votes)
|