import wandb import torch import pickle import random import torch.nn as nn import torch.nn.functional as F from operator import or_ from typing import List from functools import reduce from enum import IntFlag, auto from src.model.base import ODBinaryMetricClassifier from src.model.clip import VideoAttrExtractor from src.utility.loss import focal_loss class OpMode(IntFlag): S = auto() # spatial T = auto() # temporal def call_module(module): def fn(*args, **kwargs): return module(*args, **kwargs) return fn def get_module(module): def fn(): return module return fn class SynoBlock(nn.Module): def __init__( self, n_synos, d_model, n_head, n_patch, n_frames, ksize_t, ksize_s, t_attrs, s_k_attr, s_v_attr, op_mode, store_attrs=[], attn_record=False ): super().__init__() # parameters self.n_patch = n_patch self.t_attrs = t_attrs self.s_k_attr = s_k_attr self.s_v_attr = s_v_attr self.op_mode = op_mode if (OpMode.T in op_mode): # modules self.t_conv = self.make_2dconv( ksize_t, sum([ 1 if attr in ["out", "emb"] else n_head for attr in t_attrs ]), 1 ) self.t_proj = nn.Sequential( nn.LayerNorm(n_frames**2), nn.Linear( n_frames**2, n_frames ), nn.GELU(), nn.Linear( n_frames, n_frames**2 ) ) self.p_conv = self.make_2dconv( ksize_s, n_frames ** 2, 1 ) if (OpMode.S in op_mode): self.syno_embedding = nn.Parameter( torch.zeros(n_synos, d_model) ) # attribute storage self.store_attrs = store_attrs self.attr = {} # attention map recording self.attn_record = attn_record self.aff = None def make_2dconv(self, ksize, in_c, out_c, groups=1): conv = nn.Conv2d( in_channels=in_c, out_channels=out_c, kernel_size=ksize, stride=1, padding=ksize // 2, groups=groups, bias=True ) nn.init.normal_(conv.weight, std=0.001) nn.init.zeros_(conv.bias) return conv def pop_attr(self): ret = self.get_attr() self.attr.clear() return ret def get_attr(self): return {k: self.attr[k] for k in self.attr} def set_attr(self, **attr): self.attr = { k: attr[k] for k in attr if k in self.store_attrs } def temporal_detection(self, attrs): b, t, l, h, d = attrs['q'][:, :, 1:].shape # ignore cls token p = self.n_patch # p = l ** 0.5 affs = [] for attr in self.t_attrs: _attr = attrs[attr][:, :, 1:] # ignore cls token if (len(_attr.shape) == 4): _attr = _attr.unsqueeze(-2) _attr = _attr.permute(0, 2, 1, 3, 4) aff = torch.einsum( 'nlqhc,nlkhc->nlqkh', _attr / (_attr.size(-1) ** 0.5), _attr ) aff = aff.softmax(dim=-2) aff = aff.flatten(0, 1) # shape = (n*l,t,t,h) aff = aff.permute(0, 3, 1, 2) # shape = (n*l,h,t,t) affs.append(aff) aff = torch.cat(affs, dim=1) # shape = (n*l, 3*h, t, t) aff = self.t_conv(aff) # shape = (n*l, r, t, t) where r is number of filters aff = aff.unflatten(0, (b, p, p)).flatten(3) # shape = (n, p, p, t*t) aff = aff + self.t_proj(aff) aff = aff.permute(0, 3, 1, 2) # shape = (n, t*t, p, p) aff = self.p_conv(aff) # shape = (n, 1, p, p) y = aff.flatten(1) return dict(y=y) # shape = (n, p*p) def spatial_detection(self, attrs): b, t, l, h, d = attrs['q'][:, :, 1:].shape # ignore cls token _k = attrs[self.s_k_attr][:, :, 1:] # ignore cls token _v = attrs[self.s_v_attr][:, :, 1:] # ignore cls token # prepare query s_q = self.syno_embedding.unsqueeze(0).repeat(b, 1, 1) # shape = (b, synos, width) if (len(_k.shape) == 5): _k = _k.flatten(-2).contiguous() # match shape if (len(_v.shape) == 5): _v = _v.flatten(-2).contiguous() # match shape # =============================== # Version 1: Original Attention Module (Square Normalized) # s_aff = torch.einsum( # 'nqw,ntkw->ntqk', # s_q / (s_q.size(-1) ** 0.5), # _k # ) # s_aff = s_aff.softmax(dim=-1) # Version 2: Modified Attention Module (Cosine Similarity) s_aff = torch.einsum( 'nqw,ntkw->ntqk', s_q / (s_q.norm(dim=-1, keepdim=True) + 1e-4), _k / (_k.norm(dim=-1, keepdim=True) + 1e-4) ) s_aff = (s_aff * 100).softmax(dim=-1) # =============================== s_mix = torch.einsum( 'ntql,ntlw->ntqw', s_aff, _v ) y = s_mix.flatten(1, 2).mean(dim=1) # shape = (b,w) if self.attn_record: self.aff = s_aff return dict(s_q=s_q, y=y) def forward(self, attrs): self.pop_attr() y_t = 0 y_s = 0 if OpMode.T in self.op_mode: ret_t = self.temporal_detection(attrs) y_t = ret_t.pop('y') self.set_attr( **ret_t ) if OpMode.S in self.op_mode: ret_s = self.spatial_detection(attrs) y_s = ret_s.pop('y') self.set_attr( **ret_s ) return y_t, y_s class SynoDecoder(nn.Module): def __init__( self, encoder, num_synos, num_frames, ksize_s, ksize_t, t_attrs, s_k_attr, s_v_attr, op_mode, store_attrs=[], attn_record=False ): super().__init__() d_model = encoder.transformer.width n_head = encoder.transformer.heads n_patch = int((encoder.patch_num)**0.5) self.encoder = get_module(encoder) self.decoder_layers = nn.ModuleList([ SynoBlock( n_synos=num_synos, d_model=d_model, n_head=n_head, n_patch=n_patch, n_frames=num_frames, ksize_t=ksize_t, ksize_s=ksize_s, t_attrs=t_attrs, s_k_attr=s_k_attr, s_v_attr=s_v_attr, op_mode=op_mode, store_attrs=store_attrs, attn_record=attn_record ) for _ in range(encoder.transformer.layers) ]) self.op_mode = op_mode self.feat_t_dim = n_patch**2 self.feat_s_dim = d_model @property def spatial_dim(self): return self.feat_s_dim @property def temporal_dim(self): return self.feat_t_dim def forward(self, x): b = x.shape[0] layer_output = dict( y_t=[], y_s=[] ) # first, we prepare the encoder before the transformer layers. x = self.encoder()._prepare(x) # now, we alternate between synoptic and encoder layers for enc_blk, dec_blk in zip( self.encoder().transformer.resblocks, self.decoder_layers ): data = enc_blk(x) x = data["emb"] y_t, y_s = dec_blk(data) layer_output["y_t"].append(y_t) layer_output["y_s"].append(y_s) # last, we are done with the encoder, therefore skipping the _finalize step. # x = self.encoder()._finalize(x) # aggregate the layer outputs y_s = sum(layer_output["y_s"]) y_t = sum(layer_output["y_t"]) return y_s, y_t class SynoVideoAttrExtractor(VideoAttrExtractor): def __init__( self, # VideoAttrExtractor architecture, text_embed, pretrain=None, store_attrs=[], attn_record=False, # synoptic ksize_t=3, ksize_s=3, num_synos=1, num_frames=1, s_k_attr="k", s_v_attr="v", op_mode=(OpMode.S | OpMode.T), t_attrs=["q", "k", "v"], ): super(SynoVideoAttrExtractor, self).__init__( architecture=architecture, text_embed=text_embed, store_attrs=store_attrs, attn_record=attn_record, pretrain=pretrain ) self.decoder = SynoDecoder( encoder=self.model, num_synos=num_synos, num_frames=num_frames, t_attrs=t_attrs, ksize_t=ksize_t, ksize_s=ksize_s, s_k_attr=s_k_attr, s_v_attr=s_v_attr, op_mode=op_mode, store_attrs=store_attrs, attn_record=attn_record ) @property def spatial_dim(self): return self.decoder.spatial_dim @property def temporal_dim(self): return self.decoder.temporal_dim def forward(self, x): syno_s, syno_t = self.decoder(x=x) layer_attrs = [] for enc_blk, dec_blk in zip(self.model.transformer.resblocks, self.decoder.decoder_layers): layer_attrs.append( { **enc_blk.pop_attr(), **dec_blk.pop_attr() } ) return dict( layer_attrs=layer_attrs, syno_s=syno_s, syno_t=syno_t ) def train(self, mode=True): super().train(mode) if (mode): self.model.eval() self.decoder.train() return self class BinaryLinearClassifier(nn.Module): def __init__( self, *args, op_mode, **kargs, ): super().__init__() self.encoder = SynoVideoAttrExtractor( *args, **kargs, op_mode=op_mode ) self.op_mode = op_mode if (OpMode.S in self.op_mode): self.s_ln = nn.LayerNorm(self.encoder.spatial_dim) self.s_head = nn.Linear(self.encoder.spatial_dim, 2) if (OpMode.T in self.op_mode): self.t_ln = nn.LayerNorm(self.encoder.temporal_dim) self.t_head = nn.Linear(self.encoder.temporal_dim, 2) if (self.op_mode == (OpMode.T | OpMode.S)): self.a_head = nn.Linear( self.encoder.temporal_dim + self.encoder.spatial_dim, 2 ) @property def transform(self): return self.encoder.transform @property def n_px(self): return self.encoder.model.input_resolution def forward(self, x, *args, **kargs): results = self.encoder(x) logits_ = dict() if (OpMode.S in self.op_mode): _s = self.s_ln(results["syno_s"]) logits_s = self.s_head(_s) logits_["logits_s"] = logits_s logits_["logits"] = logits_s if (OpMode.T in self.op_mode): _t = self.t_ln(results["syno_t"]) logits_t = self.t_head(_t) logits = logits_t logits_["logits_t"] = logits_t logits_["logits"] = logits_t if (self.op_mode == (OpMode.S | OpMode.T)): logits_a = self.a_head(torch.cat([_s, _t], dim=-1)) logits = torch.log( ( logits_s.softmax(dim=-1) + logits_t.softmax(dim=-1) + logits_a.softmax(dim=-1) ) / 3 + 1e-4 ) logits_["logits_a"] = logits_a logits_["logits"] = logits return dict( **logits_, ** results ) class SynoVideoLearner(ODBinaryMetricClassifier): def __init__( self, num_synos: int = 1, num_frames: int = 1, ksize_s: int = 3, ksize_t: int = 3, t_attrs: List[str] = ["q", "k", "v"], s_k_attr: str = "k", s_v_attr: str = "v", op_mode: List[str] = ["S", "T"], architecture: str = 'ViT-B/16', text_embed: bool = False, pretrain: str = None, attn_record: bool = False, store_attrs: List[str] = [], is_focal_loss: bool = True, cls_weight: float = 10.0, label_weights: List[float] = [1, 1], ): super().__init__() self.save_hyperparameters() op_mode = reduce(or_, [OpMode[m] for m in op_mode]) params = dict( architecture=architecture, text_embed=text_embed, attn_record=attn_record, pretrain=pretrain, num_synos=num_synos, num_frames=num_frames, ksize_s=ksize_s, ksize_t=ksize_t, store_attrs=store_attrs, t_attrs=t_attrs, s_k_attr=s_k_attr, s_v_attr=s_v_attr, op_mode=op_mode ) self.model = BinaryLinearClassifier(**params) self.label_weights = torch.tensor(label_weights) self.cls_weight = cls_weight self.is_focal_loss = is_focal_loss @property def transform(self): return self.model.transform @property def n_px(self): return self.model.n_px def shared_step(self, batch, stage): x, y, z = batch["xyz"] indices = batch["indices"] dts_name = batch["dts_name"] names = batch["names"] output = self(x, **z) logits = output["logits"] loss = 0 # classification loss if (stage == "train"): logits = [output[k] for k in output if "logits_" in k] y = y.repeat(len(logits)) logits = torch.cat(logits, dim=0) if self.is_focal_loss: cls_loss = focal_loss( logits, y, gamma=4, weight=( self.label_weights.to(y.device) if stage == "train" else None ) ) else: cls_loss = nn.functional.cross_entropy( logits, y, reduction="none", weight=( self.label_weights.to(y.device) if stage == "train" else None ) ) loss += cls_loss.mean() * self.cls_weight self.log( f"{stage}/{dts_name}/loss", cls_loss.mean(), batch_size=logits.shape[0] ) else: # classification loss cls_loss = nn.functional.cross_entropy( logits, y, reduction="none", weight=None ) loss += cls_loss.mean() return { "logits": logits, "labels": y, "loss": loss, "dts_name": dts_name, "indices": indices, "output": output } class FFGSynoVideoLearner(SynoVideoLearner): def __init__( self, # ffg face_feature_path: str, face_parts: List[str] = [ "lips", "skin", "eyes", "nose" ], face_attn_attr: str = "k", syno_attn_attr: str = "s_q", ffg_temper: float = 30, ffg_weight: float = 1.5, ffg_layers: int = -1, ffg_reverse: bool = False, # generic architecture: str = 'ViT-B/16', text_embed: bool = False, pretrain: str = None, num_frames: int = 1, ksize_s: int = 3, ksize_t: int = 3, t_attrs: List[str] = ["q", "k", "v"], s_k_attr: str = "k", s_v_attr: str = "v", op_mode: List[str] = ["S", "T"], attn_record: bool = False, store_attrs: List[str] = [], is_focal_loss: bool = True, cls_weight: float = 10.0, label_weights: List[float] = [1, 1], ): assert 'S' in op_mode, "FFG must include the spatial branch for operation." self.num_face_parts = len(face_parts) self.face_attn_attr = face_attn_attr self.syno_attn_attr = syno_attn_attr self.ffg_temper = ffg_temper self.ffg_weight = ffg_weight self.ffg_layers = ffg_layers self.ffg_reverse = ffg_reverse super().__init__( num_frames=num_frames, num_synos=self.num_face_parts, ksize_s=ksize_s, ksize_t=ksize_t, op_mode=op_mode, t_attrs=t_attrs, s_k_attr=s_k_attr, s_v_attr=s_v_attr, architecture=architecture, text_embed=text_embed, attn_record=attn_record, pretrain=pretrain, store_attrs=set([*store_attrs, self.syno_attn_attr]), cls_weight=cls_weight, label_weights=label_weights, is_focal_loss=is_focal_loss ) self.save_hyperparameters() with open(face_feature_path, "rb") as f: _face_features = pickle.load(f) self.face_features = torch.stack( [ torch.stack([ _face_features[self.face_attn_attr][p][l] for p in face_parts ]) for l in range(self.model.encoder.model.transformer.layers) ] ) self.face_features = self.face_features.unsqueeze(1) for i, dec_blk in enumerate(self.model.encoder.decoder.decoder_layers): dec_blk.syno_embedding.data = self.face_features[i].squeeze(0).data.clone() def shared_step(self, batch, stage): result = super().shared_step(batch, stage) if (stage == "train"): dts_name = result["dts_name"] x = batch["xyz"][0] # face feature guided loss target_attn_attrs = torch.stack( [ attrs[self.syno_attn_attr] for attrs in result["output"]["layer_attrs"] ] ) # qs.shape = [layer,b,t,syno,patch,head] if (len(target_attn_attrs.shape) == 4): # shape = [l, b, synos, head*width] # for: out, emb pass elif (len(target_attn_attrs.shape) == 5): # shape = [l, b, synos, head,width] # for: q, k, v target_attn_attrs = target_attn_attrs.flatten(-2) elif (len(target_attn_attrs.shape) == 6): # shape = [l, b, t, synos, head, width] # for: q, k, v target_attn_attrs = target_attn_attrs.mean(2).flatten(-2) else: raise NotImplementedError() face_features = self.face_features.to( dtype=target_attn_attrs.dtype, device=target_attn_attrs.device ) if self.ffg_layers == -1: pass elif self.ffg_layers > 0: if (self.ffg_reverse): face_features = face_features[: self.ffg_layers] target_attn_attrs = target_attn_attrs[: self.ffg_layers] else: layers = ( self.model.encoder.model.transformer.layers - self.ffg_layers ) face_features = face_features[layers:] target_attn_attrs = target_attn_attrs[layers:] else: raise NotImplementedError() l, b, q = target_attn_attrs.shape[:3] face_features = face_features / face_features.norm(dim=-1, keepdim=True) target_attn_attrs = target_attn_attrs / target_attn_attrs.norm(dim=-1, keepdim=True) logits = self.ffg_temper * (target_attn_attrs @ face_features.transpose(-1, -2)) cls_sim = torch.nn.functional.cross_entropy( logits.flatten(0, 2), ( torch.arange( 0, self.num_face_parts ) .repeat((l * b)) .to(x.device) ), reduction="none" ).mean() self.log( f"{stage}/{dts_name}/syno_sim", cls_sim, batch_size=x.shape[0] ) result["loss"] += cls_sim * self.ffg_weight return result class FFESynoVideoLearner(SynoVideoLearner): def __init__( self, # ffg face_feature_path: str, face_parts: List[str] = [ "lips", "skin", "eyes", "nose" ], face_attn_attr: str = "k", syno_attn_attr: str = "s_q", ffg_temper: float = 30, ffg_weight: float = 1.5, ffg_layers: int = -1, ffg_reverse: bool = False, ffe_weight: float = 100, grad_temper: float = 100, # generic architecture: str = 'ViT-B/16', text_embed: bool = False, pretrain: str = None, num_frames: int = 1, ksize_s: int = 3, ksize_t: int = 3, t_attrs: List[str] = ["q", "k", "v"], s_k_attr: str = "k", s_v_attr: str = "v", op_mode: List[str] = ["S", "T"], attn_record: bool = True, store_attrs: List[str] = [], is_focal_loss: bool = True, cls_weight: float = 10.0, label_weights: List[float] = [1, 1], ): assert 'S' in op_mode, "FFG must include the spatial branch for operation." self.num_face_parts = len(face_parts) self.face_attn_attr = face_attn_attr self.syno_attn_attr = syno_attn_attr self.ffg_temper = ffg_temper self.ffg_weight = ffg_weight self.ffg_layers = ffg_layers self.ffg_reverse = ffg_reverse self.ffe_weight = ffe_weight self.grad_temper = grad_temper super().__init__( num_frames=num_frames, num_synos=self.num_face_parts, ksize_s=ksize_s, ksize_t=ksize_t, op_mode=op_mode, t_attrs=t_attrs, s_k_attr=s_k_attr, s_v_attr=s_v_attr, architecture=architecture, text_embed=text_embed, attn_record=attn_record, pretrain=pretrain, store_attrs=set([*store_attrs, self.syno_attn_attr]), cls_weight=cls_weight, label_weights=label_weights, is_focal_loss=is_focal_loss ) self.save_hyperparameters() with open(face_feature_path, "rb") as f: _face_features = pickle.load(f) self.face_features = torch.stack( [ torch.stack([ _face_features[self.face_attn_attr][p][l] for p in face_parts ]) for l in range(self.model.encoder.model.transformer.layers) ] ) self.face_features = self.face_features.unsqueeze(1) for i, dec_blk in enumerate(self.model.encoder.decoder.decoder_layers): dec_blk.syno_embedding.data = self.face_features[i].squeeze(0).data.clone() def shared_step(self, batch, stage): result = super().shared_step(batch, stage) if (stage == "train"): dts_name = result["dts_name"] x = batch["xyz"][0] # face feature guided loss target_attn_attrs = torch.stack( [ attrs[self.syno_attn_attr] for attrs in result["output"]["layer_attrs"] ] ) # qs.shape = [layer,b,t,syno,patch,head] if (len(target_attn_attrs.shape) == 4): # shape = [l, b, synos, head*width] # for: out, emb pass elif (len(target_attn_attrs.shape) == 5): # shape = [l, b, synos, head,width] # for: q, k, v target_attn_attrs = target_attn_attrs.flatten(-2) elif (len(target_attn_attrs.shape) == 6): # shape = [l, b, t, synos, head, width] # for: q, k, v target_attn_attrs = target_attn_attrs.mean(2).flatten(-2) else: raise NotImplementedError() face_features = self.face_features.to( dtype=target_attn_attrs.dtype, device=target_attn_attrs.device ) if self.ffg_layers == -1: pass elif self.ffg_layers > 0: if (self.ffg_reverse): face_features = face_features[: self.ffg_layers] target_attn_attrs = target_attn_attrs[: self.ffg_layers] else: layers = ( self.model.encoder.model.transformer.layers - self.ffg_layers ) face_features = face_features[layers:] target_attn_attrs = target_attn_attrs[layers:] else: raise NotImplementedError() l, b, q = target_attn_attrs.shape[:3] face_features = face_features / face_features.norm(dim=-1, keepdim=True) target_attn_attrs = target_attn_attrs / target_attn_attrs.norm(dim=-1, keepdim=True) logits = self.ffg_temper * (target_attn_attrs @ face_features.transpose(-1, -2)) cls_sim = torch.nn.functional.cross_entropy( logits.flatten(0, 2), ( torch.arange( 0, self.num_face_parts ) .repeat((l * b)) .to(x.device) ), reduction="none" ).mean() self.log( f"{stage}/{dts_name}/syno_sim", cls_sim, batch_size=x.shape[0] ) result["loss"] += cls_sim * self.ffg_weight grad_raw_grid = ( torch.stack( torch.autograd.grad( result["loss"], [layer.aff for layer in self.model.encoder.decoder.decoder_layers], grad_outputs=torch.ones_like(result["loss"]), retain_graph=True, create_graph=True ) ).abs() * self.grad_temper ) attn_grid = ( torch.stack( [layer.aff for layer in self.model.encoder.decoder.decoder_layers] ) ).detach() grad_align = torch.nn.functional.kl_div( torch.nn.functional.log_softmax(grad_raw_grid.flatten(3), dim=-1), attn_grid.flatten(3), ) self.log( f"{stage}/{dts_name}/grad_align", grad_align, batch_size=x.shape[0] ) result["loss"] += grad_align * self.ffe_weight return result if __name__ == "__main__": frames = 5 # # AttrExtractor Test model = FFGSynoVideoLearner( face_feature_path="misc/L14_real_semantic_patches_v2_2000.pickle", architecture="ViT-L/14", num_frames=frames ) model.to("cuda") results = model(torch.randn(9, frames, 3, 224, 224).to("cuda")) synos = results["logits"] synos.sum().backward() # model = GlitchBlock(n_head=12, n_patch=14, n_filt=10, ksize=3, n_frames=frames) # model.to("cuda") # logits = model({ # "q": torch.randn(1, frames, 197, 12, 64).to("cuda"), # "k": torch.randn(1, frames, 197, 12, 64).to("cuda") # }) print("done")