diff --git a/ram/inference.py b/ram/inference.py index 182efc5..d686b05 100644 --- a/ram/inference.py +++ b/ram/inference.py @@ -32,10 +32,14 @@ def inference_tag2text(image, model, input_tag="None"): def inference_ram(image, model): - with torch.no_grad(): - tags, tags_chinese = model.generate_tag(image) + # with torch.no_grad(): + # tags, tags_chinese = model.generate_tag(image) - return tags[0],tags_chinese[0] + # return tags[0],tags_chinese[0] + with torch.no_grad(): + tags, chinese_tags, tag_importance = model.generate_tag(image) + + return tags, chinese_tags, tag_importance def inference_ram_openset(image, model): diff --git a/ram/models/bert.py b/ram/models/bert.py index cb90b79..edb09a4 100644 --- a/ram/models/bert.py +++ b/ram/models/bert.py @@ -36,11 +36,11 @@ SequenceClassifierOutput, TokenClassifierOutput, ) +from transformers.pytorch_utils import apply_chunking_to_forward +from .helper import find_pruneable_heads_and_indices +from .helper2 import _prune_linear_layer from transformers.modeling_utils import ( PreTrainedModel, - apply_chunking_to_forward, - find_pruneable_heads_and_indices, - prune_linear_layer, ) from transformers.utils import logging from transformers.models.bert.configuration_bert import BertConfig @@ -310,10 +310,10 @@ def prune_heads(self, heads): ) # Prune linear layers - self.self.query = prune_linear_layer(self.self.query, index) - self.self.key = prune_linear_layer(self.self.key, index) - self.self.value = prune_linear_layer(self.self.value, index) - self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) + self.self.query = _prune_linear_layer(self.self.query, index) + self.self.key = _prune_linear_layer(self.self.key, index) + self.self.value = _prune_linear_layer(self.self.value, index) + self.output.dense = _prune_linear_layer(self.output.dense, index, dim=1) # Update hyper params and store pruned heads self.self.num_attention_heads = self.self.num_attention_heads - len(heads) diff --git a/ram/models/helper.py b/ram/models/helper.py new file mode 100644 index 0000000..bf72fb4 --- /dev/null +++ b/ram/models/helper.py @@ -0,0 +1,40 @@ +import torch + +def find_pruneable_heads_and_indices( + heads: list[int], + n_heads: int, + head_size: int, + already_pruned_heads: set[int], +) -> tuple[set[int], torch.LongTensor]: + """ + Finds the heads and the flattened indices to keep, taking already-pruned heads + into account. + + Parameters + ---------- + heads : list[int] + Head indices requested for pruning. + n_heads : int + Total number of attention heads before this pruning step. + head_size : int + Size of each attention head. + already_pruned_heads : set[int] + Heads that were pruned in earlier steps. + + Returns + ------- + tuple[set[int], torch.LongTensor] + (new_heads_to_prune, flattened_indices_to_keep) + """ + mask = torch.ones(n_heads, head_size, dtype=torch.bool) + + heads = set(heads) - already_pruned_heads + + for head in heads: + # Shift the head index left by however many smaller heads + # were already removed earlier. + shifted_head = head - sum(1 for h in already_pruned_heads if h < head) + mask[shifted_head] = False + + index = torch.arange(n_heads * head_size)[mask.view(-1)].long() + return heads, index diff --git a/ram/models/helper2.py b/ram/models/helper2.py new file mode 100644 index 0000000..2ccbc87 --- /dev/null +++ b/ram/models/helper2.py @@ -0,0 +1,92 @@ +import logging +import torch +import torch.nn as nn +from transformers import modeling_utils + +# ------------------------------------------------------------------------------ +# 1. Define Fallback Implementations +# (Mirrors logic removed in transformers v4.45+) +# ------------------------------------------------------------------------------ + +def _prune_linear_layer(layer, index, dim=0): + """Fallback implementation for prune_linear_layer""" + index = index.to(layer.weight.device) + W = layer.weight.index_select(dim, index).clone().detach() + if layer.bias is not None: + if dim == 0: + b = layer.bias.index_select(dim, index).clone().detach() + else: + b = layer.bias.clone().detach() + new_size = list(layer.weight.size()) + new_size[dim] = len(index) + new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device) + new_layer.weight.requires_grad = False + new_layer.weight.copy_(W.contiguous()) + new_layer.weight.requires_grad = True + if layer.bias is not None: + new_layer.bias.requires_grad = False + new_layer.bias.copy_(b.contiguous()) + new_layer.bias.requires_grad = True + return new_layer + +def _prune_layer(layer, index, dim=None): + """Fallback for prune_layer (generic)""" + if isinstance(layer, nn.Linear): + return _prune_linear_layer(layer, index, dim=0 if dim is None else dim) + + # Handle Conv1D if available + try: + from transformers.pytorch_utils import prune_conv1d_layer + if hasattr(modeling_utils, "Conv1D") and isinstance(layer, modeling_utils.Conv1D): + return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim) + except ImportError: + pass + + return _prune_linear_layer(layer, index, dim=0 if dim is None else dim) + +def _find_pruneable_heads_and_indices(heads, n_heads, head_size, already_pruned_heads): + """Fallback for finding heads to prune""" + mask = torch.ones(n_heads, head_size) + heads = set(heads) - already_pruned_heads + for head in heads: + head = head - sum(1 if h < head else 0 for h in already_pruned_heads) + mask[head] = 0 + mask = mask.view(-1).contiguous().eq(1) + index = torch.arange(len(mask))[mask].long() + return heads, index + +# ------------------------------------------------------------------------------ +# 2. Inject Missing Functions into transformers.modeling_utils +# ------------------------------------------------------------------------------ + +# Patch 'Conv1D' if missing +if not hasattr(modeling_utils, "Conv1D"): + try: + from transformers.pytorch_utils import Conv1D + modeling_utils.Conv1D = Conv1D + except ImportError: + pass + +# Patch 'prune_linear_layer' +if not hasattr(modeling_utils, "prune_linear_layer"): + try: + from transformers.pytorch_utils import prune_linear_layer + modeling_utils.prune_linear_layer = prune_linear_layer + except ImportError: + modeling_utils.prune_linear_layer = _prune_linear_layer + +# Patch 'prune_layer' (Critical for MeshGraphormer) +if not hasattr(modeling_utils, "prune_layer"): + modeling_utils.prune_layer = _prune_layer + +# Patch 'find_pruneable_heads_and_indices' +if not hasattr(modeling_utils, "find_pruneable_heads_and_indices"): + try: + from transformers.pytorch_utils import find_pruneable_heads_and_indices + modeling_utils.find_pruneable_heads_and_indices = find_pruneable_heads_and_indices + except ImportError: + modeling_utils.find_pruneable_heads_and_indices = _find_pruneable_heads_and_indices + +logging.info("\033[32m[Transformers Fix] Successfully injected missing pruning functions for MeshGraphormer compatibility.\033[0m") + +NODE_CLASS_MAPPINGS = {} \ No newline at end of file diff --git a/ram/models/ram.py b/ram/models/ram.py index e329897..a60721e 100644 --- a/ram/models/ram.py +++ b/ram/models/ram.py @@ -303,19 +303,59 @@ def forward(self, image, caption, image_tag, parse_tag, clip_feature): return loss_t2t, loss_tag, loss_dis - def generate_tag(self, - image, - threshold=0.68, - tag_input=None, - ): + # def generate_tag(self, + # image, + # threshold=0.68, + # tag_input=None, + # ): + # label_embed = torch.nn.functional.relu(self.wordvec_proj(self.label_embed)) + + # image_embeds = self.image_proj(self.visual_encoder(image)) + # image_atts = torch.ones(image_embeds.size()[:-1], + # dtype=torch.long).to(image.device) + + # # recognized image tags using image-tag recogntiion decoder + # image_cls_embeds = image_embeds[:, 0, :] + # image_spatial_embeds = image_embeds[:, 1:, :] + + # bs = image_spatial_embeds.shape[0] + # label_embed = label_embed.unsqueeze(0).repeat(bs, 1, 1) + # tagging_embed = self.tagging_head( + # encoder_embeds=label_embed, + # encoder_hidden_states=image_embeds, + # encoder_attention_mask=image_atts, + # return_dict=False, + # mode='tagging', + # ) + + # logits = self.fc(tagging_embed[0]).squeeze(-1) + + # targets = torch.where( + # torch.sigmoid(logits) > self.class_threshold.to(image.device), + # torch.tensor(1.0).to(image.device), + # torch.zeros(self.num_class).to(image.device)) + + # tag = targets.cpu().numpy() + # tag[:,self.delete_tag_index] = 0 + # tag_output = [] + # tag_output_chinese = [] + # for b in range(bs): + # index = np.argwhere(tag[b] == 1) + # token = self.tag_list[index].squeeze(axis=1) + # tag_output.append(' | '.join(token)) + # token_chinese = self.tag_list_chinese[index].squeeze(axis=1) + # tag_output_chinese.append(' | '.join(token_chinese)) + + + # return tag_output, tag_output_chinese + def generate_tag(self, image, threshold=0.68, tag_input=None): label_embed = torch.nn.functional.relu(self.wordvec_proj(self.label_embed)) image_embeds = self.image_proj(self.visual_encoder(image)) - image_atts = torch.ones(image_embeds.size()[:-1], - dtype=torch.long).to(image.device) + image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image.device) - # recognized image tags using image-tag recogntiion decoder + # Recognized image tags using image-tag recognition decoder image_cls_embeds = image_embeds[:, 0, :] image_spatial_embeds = image_embeds[:, 1:, :] @@ -330,25 +370,31 @@ def generate_tag(self, ) logits = self.fc(tagging_embed[0]).squeeze(-1) + probabilities = torch.sigmoid(logits) # Tính xác suất bằng sigmoid targets = torch.where( - torch.sigmoid(logits) > self.class_threshold.to(image.device), + probabilities > self.class_threshold.to(image.device), torch.tensor(1.0).to(image.device), - torch.zeros(self.num_class).to(image.device)) + torch.zeros(self.num_class).to(image.device) + ) tag = targets.cpu().numpy() - tag[:,self.delete_tag_index] = 0 + tag[:, self.delete_tag_index] = 0 tag_output = [] tag_output_chinese = [] + tag_importance = [] # Lưu độ quan trọng của các tag + for b in range(bs): index = np.argwhere(tag[b] == 1) token = self.tag_list[index].squeeze(axis=1) tag_output.append(' | '.join(token)) token_chinese = self.tag_list_chinese[index].squeeze(axis=1) tag_output_chinese.append(' | '.join(token_chinese)) + # Lưu xác suất của các tag được chọn + importance = probabilities[b][index].squeeze(axis=1).cpu().numpy() + tag_importance.append(dict(zip(token, importance))) - - return tag_output, tag_output_chinese + return tag_output, tag_output_chinese, tag_importance def generate_tag_openset(self, image, diff --git a/ram/models/utils.py b/ram/models/utils.py index 1a64765..e87fd3f 100644 --- a/ram/models/utils.py +++ b/ram/models/utils.py @@ -131,7 +131,7 @@ def init_tokenizer(text_encoder_type='bert-base-uncased'): tokenizer = BertTokenizer.from_pretrained(text_encoder_type) tokenizer.add_special_tokens({'bos_token': '[DEC]'}) tokenizer.add_special_tokens({'additional_special_tokens': ['[ENC]']}) - tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0] + tokenizer.enc_token_id = tokenizer.convert_tokens_to_ids('[ENC]') return tokenizer diff --git a/requirements.txt b/requirements.txt index 35287f4..f0ffbce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ timm==0.4.12 -transformers>=4.25.1 +transformers==4.57.1 fairscale==0.4.4 pycocoevalcap torch