Spaces:
Sleeping
Sleeping
| """Decoder implementations.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| from src.models.base import BaseDecoder | |
| class DecoderLSTM(BaseDecoder): | |
| """Image-as-first-token LSTM decoder (Show and Tell style baseline).""" | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| embed_dim: int = 256, | |
| hidden_dim: int = 512, | |
| num_layers: int = 1, | |
| dropout: float = 0.0, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) | |
| self.lstm = nn.LSTM( | |
| embed_dim, | |
| hidden_dim, | |
| num_layers=num_layers, | |
| batch_first=True, | |
| dropout=dropout if num_layers > 1 else 0.0, | |
| ) | |
| # NOTE: nn.LSTM's own `dropout` arg only applies BETWEEN stacked layers, | |
| # so it has no effect with num_layers=1 (our baseline). This separate | |
| # output_dropout applies to the LSTM's output before the final | |
| # projection, giving a real regularization effect even for a | |
| # single-layer LSTM. | |
| self.output_dropout = nn.Dropout(dropout) | |
| self.fc = nn.Linear(hidden_dim, vocab_size) | |
| def forward(self, image_embed: torch.Tensor, input_seq: torch.Tensor) -> torch.Tensor: | |
| # image_embed: (batch, embed_dim) | |
| # input_seq: (batch, seq_len) token indices (teacher-forced) | |
| word_embeds = self.embedding(input_seq) # (batch, seq_len, embed_dim) | |
| image_step = image_embed.unsqueeze(1) # (batch, 1, embed_dim) | |
| lstm_input = torch.cat([image_step, word_embeds], dim=1) # (batch, seq_len+1, embed_dim) | |
| lstm_out, _ = self.lstm(lstm_input) # (batch, seq_len+1, hidden_dim) | |
| lstm_out = self.output_dropout(lstm_out) | |
| logits = self.fc(lstm_out) # (batch, seq_len+1, vocab_size) | |
| # drop the output at the image step (index 0) -- it has no word target | |
| return logits[:, 1:, :] # (batch, seq_len, vocab_size) | |
| def generate( | |
| self, | |
| image_embed: torch.Tensor, | |
| start_idx: int, | |
| end_idx: int, | |
| max_len: int, | |
| ) -> list[int]: | |
| assert image_embed.size(0) == 1, "generate() supports batch_size=1" | |
| device = image_embed.device | |
| image_step = image_embed.unsqueeze(1) | |
| _, hidden = self.lstm(image_step) | |
| current_token = torch.tensor([[start_idx]], device=device) | |
| generated: list[int] = [] | |
| for _ in range(max_len): | |
| word_embed = self.embedding(current_token) | |
| lstm_out, hidden = self.lstm(word_embed, hidden) | |
| logits = self.fc(lstm_out.squeeze(1)) | |
| predicted_idx = int(logits.argmax(dim=-1).item()) | |
| if predicted_idx == end_idx: | |
| break | |
| generated.append(predicted_idx) | |
| current_token = torch.tensor([[predicted_idx]], device=device) | |
| return generated | |
| def generate_beam( | |
| self, | |
| image_embed: torch.Tensor, | |
| start_idx: int, | |
| end_idx: int, | |
| max_len: int, | |
| beam_width: int = 3, | |
| ) -> list[int]: | |
| """Beam search decoding for a SINGLE image (batch=1). | |
| Keeps the top `beam_width` candidate sequences at each step instead of | |
| committing to a single best word (as generate()/greedy does), which | |
| can recover from a locally-suboptimal early word choice. No retraining | |
| required -- this only changes inference-time decoding. | |
| """ | |
| assert image_embed.size(0) == 1, "generate_beam() supports batch_size=1" | |
| device = image_embed.device | |
| image_step = image_embed.unsqueeze(1) # (1, 1, embed_dim) | |
| _, init_hidden = self.lstm(image_step) | |
| # each beam: (token_ids, cumulative_log_prob, hidden_state, finished) | |
| beams = [([start_idx], 0.0, init_hidden, False)] | |
| for _ in range(max_len): | |
| all_candidates = [] | |
| for tokens, log_prob, hidden, finished in beams: | |
| if finished: | |
| all_candidates.append((tokens, log_prob, hidden, finished)) | |
| continue | |
| last_token = torch.tensor([[tokens[-1]]], device=device) | |
| word_embed = self.embedding(last_token) | |
| lstm_out, new_hidden = self.lstm(word_embed, hidden) | |
| logits = self.fc(lstm_out.squeeze(1)) | |
| log_probs = torch.log_softmax(logits, dim=-1).squeeze(0) | |
| top_log_probs, top_indices = log_probs.topk(beam_width) | |
| for lp, idx in zip(top_log_probs.tolist(), top_indices.tolist()): | |
| new_tokens = tokens + [idx] | |
| new_finished = idx == end_idx | |
| all_candidates.append((new_tokens, log_prob + lp, new_hidden, new_finished)) | |
| all_candidates.sort(key=lambda c: c[1], reverse=True) | |
| beams = all_candidates[:beam_width] | |
| if all(finished for _, _, _, finished in beams): | |
| break | |
| best_tokens, _, _, _ = max(beams, key=lambda c: c[1]) | |
| result = best_tokens[1:] | |
| if result and result[-1] == end_idx: | |
| result = result[:-1] | |
| return result | |
| class BahdanauAttention(nn.Module): | |
| """Additive (Bahdanau-style) attention over spatial image regions. | |
| At each decoder step, computes a weighted combination of the encoder's | |
| spatial feature grid, where the weights depend on the current decoder | |
| hidden state -- this is what lets the decoder "look at" different image | |
| regions for different words, instead of relying on one static vector. | |
| """ | |
| def __init__(self, embed_dim: int, hidden_dim: int, attention_dim: int): | |
| super().__init__() | |
| self.encoder_att = nn.Linear(embed_dim, attention_dim) | |
| self.decoder_att = nn.Linear(hidden_dim, attention_dim) | |
| self.full_att = nn.Linear(attention_dim, 1) | |
| self.relu = nn.ReLU() | |
| self.softmax = nn.Softmax(dim=1) | |
| def forward(self, encoder_out: torch.Tensor, decoder_hidden: torch.Tensor): | |
| att1 = self.encoder_att(encoder_out) | |
| att2 = self.decoder_att(decoder_hidden).unsqueeze(1) | |
| att = self.full_att(self.relu(att1 + att2)).squeeze(2) | |
| alpha = self.softmax(att) | |
| context = (encoder_out * alpha.unsqueeze(2)).sum(dim=1) | |
| return context, alpha | |
| class DecoderAttentionLSTM(BaseDecoder): | |
| """LSTM decoder with Bahdanau-style attention over spatial image features. | |
| Unlike DecoderLSTM (which injects one static image vector as the first | |
| LSTM step), this decoder recomputes a fresh, weighted combination of | |
| image regions at EVERY generation step, using an LSTMCell in an explicit | |
| loop (required since each step's attention depends on that step's hidden | |
| state -- unlike DecoderLSTM, the whole sequence can't be processed in one | |
| nn.LSTM call). | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| embed_dim: int = 256, | |
| hidden_dim: int = 512, | |
| attention_dim: int = 256, | |
| dropout: float = 0.5, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.embed_dim = embed_dim | |
| self.hidden_dim = hidden_dim | |
| self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) | |
| self.attention = BahdanauAttention(embed_dim, hidden_dim, attention_dim) | |
| self.lstm_cell = nn.LSTMCell(embed_dim + embed_dim, hidden_dim) | |
| self.init_h = nn.Linear(embed_dim, hidden_dim) | |
| self.init_c = nn.Linear(embed_dim, hidden_dim) | |
| self.f_beta = nn.Linear(hidden_dim, embed_dim) | |
| self.sigmoid = nn.Sigmoid() | |
| self.dropout = nn.Dropout(dropout) | |
| self.fc = nn.Linear(hidden_dim, vocab_size) | |
| def _init_hidden_state(self, encoder_out: torch.Tensor): | |
| mean_encoder_out = encoder_out.mean(dim=1) | |
| h = self.init_h(mean_encoder_out) | |
| c = self.init_c(mean_encoder_out) | |
| return h, c | |
| def forward(self, image_embed: torch.Tensor, input_seq: torch.Tensor) -> torch.Tensor: | |
| batch_size = image_embed.size(0) | |
| seq_len = input_seq.size(1) | |
| vocab_size = self.fc.out_features | |
| device = image_embed.device | |
| embeddings = self.embedding(input_seq) | |
| h, c = self._init_hidden_state(image_embed) | |
| outputs = torch.zeros(batch_size, seq_len, vocab_size, device=device) | |
| for t in range(seq_len): | |
| context, _ = self.attention(image_embed, h) | |
| gate = self.sigmoid(self.f_beta(h)) | |
| context = gate * context | |
| lstm_input = torch.cat([embeddings[:, t, :], context], dim=1) | |
| h, c = self.lstm_cell(lstm_input, (h, c)) | |
| preds = self.fc(self.dropout(h)) | |
| outputs[:, t, :] = preds | |
| return outputs | |
| def generate( | |
| self, | |
| image_embed: torch.Tensor, | |
| start_idx: int, | |
| end_idx: int, | |
| max_len: int, | |
| ) -> list[int]: | |
| assert image_embed.size(0) == 1, "generate() supports batch_size=1" | |
| device = image_embed.device | |
| h, c = self._init_hidden_state(image_embed) | |
| current_token = torch.tensor([start_idx], device=device) | |
| generated: list[int] = [] | |
| for _ in range(max_len): | |
| word_embed = self.embedding(current_token) | |
| context, _ = self.attention(image_embed, h) | |
| gate = self.sigmoid(self.f_beta(h)) | |
| context = gate * context | |
| lstm_input = torch.cat([word_embed, context], dim=1) | |
| h, c = self.lstm_cell(lstm_input, (h, c)) | |
| logits = self.fc(h) | |
| predicted_idx = int(logits.argmax(dim=-1).item()) | |
| if predicted_idx == end_idx: | |
| break | |
| generated.append(predicted_idx) | |
| current_token = torch.tensor([predicted_idx], device=device) | |
| return generated | |
| def generate_beam( | |
| self, | |
| image_embed: torch.Tensor, | |
| start_idx: int, | |
| end_idx: int, | |
| max_len: int, | |
| beam_width: int = 3, | |
| ) -> list[int]: | |
| """Beam search decoding for a SINGLE image (batch=1). | |
| Same beam-search principle as DecoderLSTM.generate_beam(), adapted for | |
| the LSTMCell + per-step attention loop: each beam carries its own | |
| (h, c) hidden state, since attention (and therefore the next word | |
| distribution) depends on each beam's own hidden state, not a shared one. | |
| """ | |
| assert image_embed.size(0) == 1, "generate_beam() supports batch_size=1" | |
| device = image_embed.device | |
| init_h, init_c = self._init_hidden_state(image_embed) | |
| # each beam: (token_ids, cumulative_log_prob, h, c, finished) | |
| beams = [([start_idx], 0.0, init_h, init_c, False)] | |
| for _ in range(max_len): | |
| all_candidates = [] | |
| for tokens, log_prob, h, c, finished in beams: | |
| if finished: | |
| all_candidates.append((tokens, log_prob, h, c, finished)) | |
| continue | |
| last_token = torch.tensor([tokens[-1]], device=device) | |
| word_embed = self.embedding(last_token) | |
| context, _ = self.attention(image_embed, h) | |
| gate = self.sigmoid(self.f_beta(h)) | |
| context = gate * context | |
| lstm_input = torch.cat([word_embed, context], dim=1) | |
| new_h, new_c = self.lstm_cell(lstm_input, (h, c)) | |
| logits = self.fc(new_h) | |
| log_probs = torch.log_softmax(logits, dim=-1).squeeze(0) | |
| top_log_probs, top_indices = log_probs.topk(beam_width) | |
| for lp, idx in zip(top_log_probs.tolist(), top_indices.tolist()): | |
| new_tokens = tokens + [idx] | |
| new_finished = idx == end_idx | |
| all_candidates.append((new_tokens, log_prob + lp, new_h, new_c, new_finished)) | |
| all_candidates.sort(key=lambda cand: cand[1], reverse=True) | |
| beams = all_candidates[:beam_width] | |
| if all(finished for _, _, _, _, finished in beams): | |
| break | |
| best_tokens, _, _, _, _ = max(beams, key=lambda cand: cand[1]) | |
| result = best_tokens[1:] | |
| if result and result[-1] == end_idx: | |
| result = result[:-1] | |
| return result | |