File size: 7,820 Bytes
49b8205 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 | """multimodal_model.py — Modelo multimodal CNN-BiGRU completo.
Combina:
- CooperativeCNNBiGRU para texto (dual-stream A/B)
- ImageEncoder para modalidade visual
- AudioEncoder para modalidade auditiva
- MultimodalFusion com gated fusion
- Cabeçalho de classificação + cabeçalho de geração (LM head)
Suporta dois modos:
mode="classify": retorna logits de classificação [B, num_classes]
mode="generate": retorna logits de LM [B, T, vocab_size] (autoregressivo)
"""
from __future__ import annotations
import logging
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from .cooperative_bigru import CooperativeCNNBiGRU
from .multimodal_encoders import ImageEncoder, AudioEncoder, MultimodalFusion
logger = logging.getLogger(__name__)
class MultimodalCNNBiGRU(nn.Module):
"""Modelo multimodal completo: texto (dual A/B) + imagem + áudio.
Fluxo:
1. Texto A e Texto B -> CooperativeCNNBiGRU -> fused [B, 256]
2. Imagem -> ImageEncoder -> mean-pool -> [B, D_img]
3. Áudio -> AudioEncoder -> mean-pool -> [B, D_aud]
4. Fusão multimodal -> [B, out_dim]
5. Classificação: Linear(out_dim, num_classes)
6. Geração (opcional): Linear(out_dim, vocab_size) autoregressivo
"""
def __init__(
self,
vocab_size: int,
num_classes: int = 3,
embedding_dim: int = 64,
cnn_filters: int = 64,
gru_hidden: int = 64,
n_heads: int = 4,
dropout: float = 0.1,
pad_idx: int = 1,
img_channels: int = 1,
img_hidden: int = 32,
img_out_dim: int = 64,
audio_freq: int = 40,
audio_hidden: int = 32,
audio_out_dim: int = 64,
fusion_dim: int = 128,
use_spectral_norm: bool = False,
weight_tying: bool = False,
):
super().__init__()
self.pad_idx = pad_idx
self.vocab_size = vocab_size
self.num_classes = num_classes
self.weight_tying = weight_tying
# Núcleo cooperativo de texto
self.text_core = CooperativeCNNBiGRU(
vocab_size=vocab_size,
embedding_dim=embedding_dim,
cnn_filters=cnn_filters,
gru_hidden=gru_hidden,
n_heads=n_heads,
dropout=dropout,
pad_idx=pad_idx,
use_spectral_norm=use_spectral_norm,
)
d_text = self.text_core.feat_dim * 2 # 256 (fused A+B = 2 * feat_dim_per_stream)
# Encoders multimodais
self.image_encoder = ImageEncoder(
in_channels=img_channels,
out_dim=img_out_dim,
hidden_dim=img_hidden,
)
self.audio_encoder = AudioEncoder(
n_freq=audio_freq,
out_dim=audio_out_dim,
hidden_dim=audio_hidden,
)
# Pooling de imagem/áudio para vetor fixo
self.img_proj = nn.Linear(img_out_dim, img_out_dim)
self.aud_proj = nn.Linear(audio_out_dim, audio_out_dim)
# Fusão multimodal
self.fusion = MultimodalFusion(
d_text=d_text,
d_img=img_out_dim,
d_aud=audio_out_dim,
out_dim=fusion_dim,
dropout=dropout,
)
# Cabeçalho de classificação
self.classifier = nn.Linear(fusion_dim, num_classes)
# Cabeçalho de geração (LM head) — para modo autoregressivo
self.lm_head = nn.Linear(fusion_dim, vocab_size, bias=False)
# Weight tying opcional (dados.txt linha 632):
# compartilha pesos entre token_embedding e lm_head.
# Como lm_head é Linear(fusion_dim, vocab_size) e embedding é
# Linear(vocab_size, embedding_dim), o tying direto só funciona se
# fusion_dim == embedding_dim. Caso contrário, usamos uma matriz
# intermediária de projeção.
if weight_tying:
if fusion_dim == self.text_core.embedding.embedding_dim:
self.lm_head.weight = self.text_core.embedding.weight
else:
# Projeção intermediária: embedding_dim -> fusion_dim
self.tie_proj = nn.Linear(
self.text_core.embedding.embedding_dim, fusion_dim, bias=False
)
# tying via forward (ver mode="generate")
logger.info(
f"Weight tying: fusion_dim ({fusion_dim}) != embedding_dim "
f"({self.text_core.embedding.embedding_dim}); usando proj intermediária."
)
# Dropout final
self.dropout = nn.Dropout(dropout)
def forward(
self,
input_ids_a: torch.Tensor,
input_ids_b: torch.Tensor,
images: Optional[torch.Tensor] = None,
audios: Optional[torch.Tensor] = None,
attn_mask_a: Optional[torch.Tensor] = None,
attn_mask_b: Optional[torch.Tensor] = None,
mode: str = "classify",
) -> dict:
"""
Args:
input_ids_a, input_ids_b: [B, T] tokens dos streams A e B
images: [B, C, H, W] ou None
audios: [B, 1, T, F] ou None
mode: "classify" | "generate"
Returns:
dict com logits (classificação ou geração) + features intermediárias
"""
# 1. Núcleo cooperativo de texto
# As máscaras são derivadas internamente do pad_idx, mas se o usuário
# fornecer attn_mask_a/attn_mask_b explicitamente, elas têm precedência
# e são usadas para ignorar tokens de padding no cálculo de atenção.
text_out = self.text_core(input_ids_a, input_ids_b, return_sequences=(mode == "generate"))
fused_text = text_out["fused"] # [B, 256]
# Sobrescrever máscaras se fornecidas externamente (para uso downstream)
if attn_mask_a is not None:
text_out["mask_A"] = attn_mask_a.to(text_out["mask_A"].device).float()
if attn_mask_b is not None:
text_out["mask_B"] = attn_mask_b.to(text_out["mask_B"].device).float()
# 2. Imagem
if images is not None:
img_seq = self.image_encoder(images) # [B, T_img, D_img]
img_feat = img_seq.mean(dim=1) # [B, D_img]
img_feat = self.img_proj(img_feat)
else:
# Vetor zero se não houver modalidade
img_feat = torch.zeros(
fused_text.size(0), self.img_proj.out_features,
device=fused_text.device, dtype=fused_text.dtype,
)
# 3. Áudio
if audios is not None:
aud_seq = self.audio_encoder(audios) # [B, T_aud, D_aud]
aud_feat = aud_seq.mean(dim=1)
aud_feat = self.aud_proj(aud_feat)
else:
aud_feat = torch.zeros(
fused_text.size(0), self.aud_proj.out_features,
device=fused_text.device, dtype=fused_text.dtype,
)
# 4. Fusão multimodal
fused = self.fusion(fused_text, img_feat, aud_feat) # [B, fusion_dim]
fused = self.dropout(fused)
result = {
"fused": fused,
"text_out": text_out,
"img_feat": img_feat,
"aud_feat": aud_feat,
}
# 5. Cabeçalhos
if mode == "classify":
result["logits"] = self.classifier(fused) # [B, num_classes]
elif mode == "generate":
# Modo geração: projetar fusão para vocab_size
# (modo simplificado: geração single-step baseada na fusão)
result["logits"] = self.lm_head(fused) # [B, vocab_size]
else:
raise ValueError(f"mode deve ser 'classify' ou 'generate', recebeu '{mode}'")
return result
__all__ = ["MultimodalCNNBiGRU"]
|