from torch import Tensor import torch from torch import nn import torch.utils.checkpoint import contextlib import torchvision from einops import rearrange import math from Uni_Sign.stgcn_layers import Graph, get_stgcn_chain from Uni_Sign.deformable_attention_2d import DeformableAttention2D from transformers import MT5ForConditionalGeneration, T5Tokenizer, MT5Config import warnings mt5_path = r"./Uni_Sign/unisign_model" def _no_grad_trunc_normal_(tensor, mean, std, a, b): # Cut & paste from PyTorch official master until it's in a few official releases - RW # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf def norm_cdf(x): # Computes standard normal cumulative distribution function return (1. + math.erf(x / math.sqrt(2.))) / 2. if (mean < a - 2 * std) or (mean > b + 2 * std): warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " "The distribution of values may be incorrect.", stacklevel=2) with torch.no_grad(): # Values are generated by using a truncated uniform distribution and # then using the inverse CDF for the normal distribution. # Get upper and lower cdf values l = norm_cdf((a - mean) / std) u = norm_cdf((b - mean) / std) # Uniformly fill tensor with values from [l, u], then translate to # [2l-1, 2u-1]. tensor.uniform_(2 * l - 1, 2 * u - 1) # Use inverse cdf transform for normal distribution to get truncated # standard normal tensor.erfinv_() # Transform to proper mean, std tensor.mul_(std * math.sqrt(2.)) tensor.add_(mean) # Clamp to ensure it's in the proper range tensor.clamp_(min=a, max=b) return tensor def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): # type: (Tensor, float, float, float, float) -> Tensor r"""Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for generating the random values works best when :math:`a \leq \text{mean} \leq b`. Args: tensor: an n-dimensional `torch.Tensor` mean: the mean of the normal distribution std: the standard deviation of the normal distribution a: the minimum cutoff value b: the maximum cutoff value Examples: >>> w = torch.empty(3, 5) >>> nn.init.trunc_normal_(w) """ return _no_grad_trunc_normal_(tensor, mean, std, a, b) class Uni_Sign(nn.Module): def __init__(self, args): super(Uni_Sign, self).__init__() self.args = args self.modes = ['body', 'left', 'right', 'face_all'] self.graph, A = {}, [] # project (x,y,score) to hidden dim hidden_dim = args.hidden_dim self.proj_linear = nn.ModuleDict() for mode in self.modes: graph_layout = f'{args.layout}_ytasl_{mode}' if self.args.dataset in ["YTASL", "Isharah"] else f'{args.layout}_{mode}' self.graph[mode] = Graph(layout=graph_layout, strategy='distance', max_hop=1) A.append(torch.tensor(self.graph[mode].A, dtype=torch.float32, requires_grad=False)) self.proj_linear[mode] = nn.Linear(3, 64) self.gcn_modules = nn.ModuleDict() self.fusion_gcn_modules = nn.ModuleDict() spatial_kernel_size = A[0].size(0) for index, mode in enumerate(self.modes): self.gcn_modules[mode], final_dim = get_stgcn_chain(64, 'spatial', (1, spatial_kernel_size), A[index].clone(), adaptive=not self.args.no_adaptive_gcn) self.fusion_gcn_modules[mode], _ = get_stgcn_chain(final_dim, 'temporal', (5, spatial_kernel_size), A[index].clone(), adaptive=not self.args.no_adaptive_gcn) self.gcn_modules['left'] = self.gcn_modules['right'] self.fusion_gcn_modules['left'] = self.fusion_gcn_modules['right'] self.proj_linear['left'] = self.proj_linear['right'] self.part_para = nn.Parameter(torch.zeros(hidden_dim*len(self.modes))) self.pose_proj = nn.Linear(256*4, 768) self.apply(self._init_weights) if self.args.dataset == "Isharah": self.lang = 'Arabic' elif "CSL" in self.args.dataset: self.lang = 'Chinese' else: self.lang = 'English' if self.args.rgb_support: self.rgb_support_backbone = torch.nn.Sequential(*list(torchvision.models.efficientnet_b0(pretrained=True).children())[:-2]) self.rgb_proj = nn.Conv2d(1280, hidden_dim, kernel_size=1) self.fusion_pose_rgb_linear = nn.Linear(hidden_dim, hidden_dim) # PGF self.fusion_pose_rgb_DA = DeformableAttention2D( dim = hidden_dim, # feature dimensions dim_head = 32, # dimension per head heads = 8, # attention heads dropout = 0., # dropout downsample_factor = 1, # downsample factor (r in paper) offset_scale = None, # scale of offset, maximum offset offset_groups = None, # number of offset groups, should be multiple of heads offset_kernel_size = 1, # offset kernel size ) self.fusion_gate = nn.Sequential(nn.Conv1d(hidden_dim*2, hidden_dim, 1), nn.GELU(), nn.Conv1d(hidden_dim, 1, 1), nn.Tanh(), nn.ReLU(), ) for layer in self.fusion_gate: try: if isinDataLoaderance(layer, nn.Conv1d): nn.init.constant_(layer.weight, 0) nn.init.constant_(layer.bias, 0) except: print("NOT IMPLEMENTED...") # Načte pouze strukturu architektury z config.json mt5_config = MT5Config.from_pretrained(mt5_path) # Vytvoří model s prázdnými vahami, které hned v dalším kroku přepíšeme self.mt5_model = MT5ForConditionalGeneration(mt5_config) self.mt5_tokenizer = T5Tokenizer.from_pretrained(mt5_path, legacy=False) self.n_registers = args.n_registers self.register_position = args.register_position self.d_model = self.mt5_model.config.d_model # should be 768 if self.n_registers > 0: self.register_tokens = nn.Parameter(torch.zeros(self.n_registers, self.d_model)) # init like other embeddings trunc_normal_(self.register_tokens, std=0.02) else: self.register_tokens = None def _init_weights(self, m): if isinstance(m, nn.Linear): trunc_normal_(m.weight, std=.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def maybe_autocast(self, dtype=torch.float32): # if on cpu, don't use autocast # if on gpu, use autocast with dtype if provided, otherwise use torch.float16 # enable_autocast = self.device != torch.device("cpu") enable_autocast = True if enable_autocast: return torch.cuda.amp.autocast(dtype=dtype) else: return contextlib.nullcontext() def gather_feat_pose_rgb(self, gcn_feat, rgb_feat, indices, rgb_len, pose_init): b, c, T, n = gcn_feat.shape assert rgb_feat.shape[0] == indices.shape[0] rgb_feat = self.rgb_proj(rgb_feat) assert len(rgb_len) == b start = 0 for batch in range(b): index = indices[start:start + rgb_len[batch]].to(torch.long) # ignore some invalid rgb clip if rgb_len[batch] == 1 and -1 in index: start = start + rgb_len[batch] continue # index selection gcn_feat_selected = gcn_feat[batch, :, index] rgb_feat_selected = rgb_feat[start:start + rgb_len[batch]] pose_init_selected = pose_init[start:start + rgb_len[batch]] gcn_feat_selected = rearrange(gcn_feat_selected, 'c t n -> t c n') pose_init_selected = rearrange(pose_init_selected, 't n c -> t c n') # PGF forward with self.maybe_autocast(): fused_transposed = self.fusion_pose_rgb_DA(pose_feat=gcn_feat_selected, rgb_feat=rgb_feat_selected, pose_init=pose_init_selected, ) fused_transposed = fused_transposed.to(gcn_feat.dtype) gate_feature = torch.concat([fused_transposed, gcn_feat_selected,], dim=-2) gate_score = self.fusion_gate(gate_feature) fused_transposed_post = (gate_score) * fused_transposed + (1 - gate_score) * gcn_feat_selected gcn_feat = gcn_feat.clone() fused_transposed_post = rearrange(fused_transposed_post, 't c n -> c t n') # replace gcn feature gcn_feat[batch, :, index] = fused_transposed_post start = start + rgb_len[batch] assert start == rgb_feat.shape[0] return gcn_feat def forward(self, src_input, tgt_input): # RGB branch forward if self.args.rgb_support: rgb_support_dict = {} for index_key, rgb_key in zip(['left_sampled_indices', 'right_sampled_indices'], ['left_hands', 'right_hands']): rgb_feat = self.rgb_support_backbone(src_input[rgb_key]) rgb_support_dict[index_key] = src_input[index_key] rgb_support_dict[rgb_key] = rgb_feat # Pose branch forward features = [] body_feat = None for part in self.modes: # project position to hidden dim proj_feat = self.proj_linear[part](src_input[part]).permute(0,3,1,2) #B,C,T,V # spatial gcn forward gcn_feat = self.gcn_modules[part](proj_feat) if part == 'body': body_feat = gcn_feat else: assert not body_feat is None if part == 'left': # Pose RGB fusion if self.args.rgb_support: gcn_feat = self.gather_feat_pose_rgb(gcn_feat, rgb_support_dict[f'{part}_hands'], rgb_support_dict[f'{part}_sampled_indices'], src_input[f'{part}_rgb_len'], src_input[f'{part}_skeletons_norm'], ) gcn_feat = gcn_feat + body_feat[..., -2][...,None].detach() elif part == 'right': # Pose RGB fusion if self.args.rgb_support: gcn_feat = self.gather_feat_pose_rgb(gcn_feat, rgb_support_dict[f'{part}_hands'], rgb_support_dict[f'{part}_sampled_indices'], src_input[f'{part}_rgb_len'], src_input[f'{part}_skeletons_norm'], ) gcn_feat = gcn_feat + body_feat[..., -1][...,None].detach() elif part == 'face_all': gcn_feat = gcn_feat + body_feat[..., 0][...,None].detach() else: raise NotImplementedError # temporal gcn forward gcn_feat = self.fusion_gcn_modules[part](gcn_feat) #B,C,T,V pool_feat = gcn_feat.mean(-1).transpose(1,2) #B,T,C features.append(pool_feat) # concat sub-pose feature across token dimension inputs_embeds = torch.cat(features, dim=-1) + self.part_para inputs_embeds = self.pose_proj(inputs_embeds) prefix_token = self.mt5_tokenizer( [f"Translate sign language video to {self.lang}: "] * len(tgt_input["gt_sentence"]), padding="longest", truncation=True, return_tensors="pt", ).to(inputs_embeds.device) prefix_embeds = self.mt5_model.encoder.embed_tokens(prefix_token['input_ids']) if self.n_registers > 0: B = inputs_embeds.size(0) # expand registers for batch register_embeds = self.register_tokens.unsqueeze(0).expand(B, -1, -1) # shape: (B, 4, 768) register_mask = torch.ones((B, self.n_registers), device=inputs_embeds.device,dtype=prefix_token['attention_mask'].dtype) if self.register_position == 'before_all': # prepend order: [registers | prefix | pose_tokens] inputs_embeds = torch.cat([register_embeds, prefix_embeds, inputs_embeds], dim=1) attention_mask = torch.cat([register_mask, prefix_token['attention_mask'], src_input['attention_mask']], dim=1) elif self.register_position == 'after_prefix' or self.register_position == 'before_pose': # prepend order: [prefix | registers | pose_tokens] inputs_embeds = torch.cat([prefix_embeds, register_embeds, inputs_embeds], dim=1) attention_mask = torch.cat([prefix_token['attention_mask'], register_mask, src_input['attention_mask']], dim=1) elif self.register_position == 'after_valid_pose': # prepend order: [prefix | valid_pose_tokens | registers | padded_pose_tokens] inputs_list = [] mask_list = [] for b in range(B): valid_len = int(src_input['attention_mask'][b].sum().item()) pose_valid = inputs_embeds[b, :valid_len] pose_pad = inputs_embeds[b, valid_len:] emb = torch.cat( [prefix_embeds[b], pose_valid, register_embeds[b], pose_pad], dim=0 ) m = torch.cat( [prefix_token['attention_mask'][b], src_input['attention_mask'][b, :valid_len], register_mask[b], src_input['attention_mask'][b, valid_len:]], dim=0 ) inputs_list.append(emb) mask_list.append(m) inputs_embeds = torch.stack(inputs_list, dim=0) attention_mask = torch.stack(mask_list, dim=0) elif self.register_position == 'after_all': # prepend order: [prefix | pose_tokens | registers] inputs_embeds = torch.cat([prefix_embeds, inputs_embeds, register_embeds], dim=1) attention_mask = torch.cat([prefix_token['attention_mask'], src_input['attention_mask'], register_mask], dim=1) else: # prepend order: [prefix | pose_tokens] inputs_embeds = torch.cat([prefix_embeds, inputs_embeds], dim=1) attention_mask = torch.cat([prefix_token['attention_mask'], src_input['attention_mask']], dim=1) tgt_input_tokenizer = self.mt5_tokenizer(tgt_input['gt_sentence'], return_tensors="pt", padding=True, truncation=True, max_length=50) labels = tgt_input_tokenizer['input_ids'] labels[labels == self.mt5_tokenizer.pad_token_id] = -100 out = self.mt5_model(inputs_embeds = inputs_embeds, attention_mask = attention_mask, labels = labels.to(inputs_embeds.device), return_dict = True, ) label = labels.reshape(-1) out_logits = out['logits'] logits = out_logits.reshape(-1,out_logits.shape[-1]) loss_fct = torch.nn.CrossEntropyLoss(label_smoothing=self.args.label_smoothing, ignore_index=-100) loss = loss_fct(logits, label.to(out_logits.device, non_blocking=True)) stack_out = { # use for inference 'inputs_embeds':inputs_embeds, 'attention_mask':attention_mask, 'loss':loss, } return stack_out @torch.no_grad() def generate(self,pre_compute_item,max_new_tokens,num_beams): inputs_embeds = pre_compute_item['inputs_embeds'] attention_mask = pre_compute_item['attention_mask'] out = self.mt5_model.generate(inputs_embeds = inputs_embeds, attention_mask = attention_mask, max_new_tokens=max_new_tokens, num_beams = num_beams, ) return out def get_requires_grad_dict(model): param_requires_grad = {name: True for name, param in model.named_parameters()} param_requires_grad_right = {} for key in param_requires_grad.keys(): if 'left' in key: param_requires_grad_right[key.replace("left", 'right')] = param_requires_grad[key] param_requires_grad = {**param_requires_grad, **param_requires_grad_right} params_to_update = {k: v for k, v in model.state_dict().items() if param_requires_grad.get(k, True)} return params_to_update