question:原始模型：import math import logging from functools import partial from collections import OrderedDict from copy import deepcopy import torch import torch.nn as nn import torch.nn.functional as F from timm.models.layers import to_2tuple from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise from .utils import combine_tokens, recover_tokens from .vit import VisionTransformer from ..layers.attn_blocks import CEBlock _logger = logging.getLogger(__name__) class VisionTransformerCE(VisionTransformer): """ Vision Transformer with candidate elimination (CE) module A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 Includes distillation token & head support for `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877 """ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False, drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None, weight_init='', ce_loc=None, ce_keep_ratio=None): """ Args: img_size (int, tuple): input image size patch_size (int, tuple): patch size in_chans (int): number of input channels num_classes (int): number of classes for classification head embed_dim (int): embedding dimension depth (int): depth of transformer num_heads (int): number of attention heads mlp_ratio (int): ratio of mlp hidden dim to embedding dim qkv_bias (bool): enable bias for qkv if True representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set distilled (bool): model includes a distillation token and head as in DeiT models drop_rate (float): dropout rate attn_drop_rate (float): attention dropout rate drop_path_rate (float): stochastic depth rate embed_layer (nn.Module): patch embedding layer norm_layer: (nn.Module): normalization layer weight_init: (str): weight init scheme """ # super().__init__() super().__init__() if isinstance(img_size, tuple): self.img_size = img_size else: self.img_size = to_2tuple(img_size) self.patch_size = patch_size self.in_chans = in_chans self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models self.num_tokens = 2 if distilled else 1 norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4) # self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4) # self.pos_embed_event_z = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=3, stride=1) # attn = CrossAttn(768, 4, 3072, 0.1, 'relu') # self.cross_attn = Iter_attn(attn, 2) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule blocks = [] ce_index = 0 self.ce_loc = ce_loc for i in range(depth): ce_keep_ratio_i = 1.0 if ce_loc is not None and i in ce_loc: ce_keep_ratio_i = ce_keep_ratio[ce_index] ce_index += 1 blocks.append( CEBlock( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, keep_ratio_search=ce_keep_ratio_i) ) self.blocks = nn.Sequential(*blocks) self.norm = norm_layer(embed_dim) self.init_weights(weight_init) def forward_features(self, z, x, event_z, event_x, mask_z=None, mask_x=None, ce_template_mask=None, ce_keep_rate=None, return_last_attn=False ): B, H, W = x.shape[0], x.shape[2], x.shape[3] event_z = self.pos_embed_event(event_z) # [:,:,:,:1000] event_x = self.pos_embed_event(event_x) # B 768 1024 x = self.patch_embed(x) z = self.patch_embed(z) event_z += self.pos_embed_z event_x += self.pos_embed_x z += self.pos_embed_z x += self.pos_embed_x # attention mask handling # B, H, W if mask_z is not None and mask_x is not None: mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_z = mask_z.flatten(1).unsqueeze(-1) mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_x = mask_x.flatten(1).unsqueeze(-1) mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode) mask_x = mask_x.squeeze(-1) if self.add_cls_token: cls_tokens = self.cls_token.expand(B, -1, -1) cls_tokens = cls_tokens + self.cls_pos_embed if self.add_sep_seg: x += self.search_segment_pos_embed z += self.template_segment_pos_embed x = combine_tokens(z, event_z, x, event_x, mode=self.cat_mode) # 64+64+256+256=640 # x = combine_tokens(z, x, event_z, event_x, mode=self.cat_mode) # 64+64+256+256=640 if self.add_cls_token: x = torch.cat([cls_tokens, x], dim=1) x = self.pos_drop(x) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device) global_index_t = global_index_t.repeat(B, 1) global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device) global_index_s = global_index_s.repeat(B, 1) removed_indexes_s = [] for i, blk in enumerate(self.blocks): x, global_index_t, global_index_s, removed_index_s, attn = blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate) if self.ce_loc is not None and i in self.ce_loc: removed_indexes_s.append(removed_index_s) x = self.norm(x) lens_x_new = global_index_s.shape[1] lens_z_new = global_index_t.shape[1] z = x[:, :lens_z_new*2] x = x[:, lens_z_new*2:] if removed_indexes_s and removed_indexes_s[0] is not None: removed_indexes_cat = torch.cat(removed_indexes_s, dim=1) pruned_lens_x = lens_x - lens_x_new pad_x = torch.zeros([B, pruned_lens_x, x.shape[2]], device=x.device) x = torch.cat([x, pad_x], dim=1) index_all = torch.cat([global_index_s, removed_indexes_cat], dim=1) # recover original token order C = x.shape[-1] x = torch.zeros_like(x).scatter_(dim=1, index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64), src=x) x = recover_tokens(x, lens_z_new, lens_x, mode=self.cat_mode) x = x[:, :lens_x] # RGB head x = torch.cat([event_x, x], dim=1) # x = x[:, lens_x//2:] # event head # x = torch.cat([z, x], dim=1) # re-concatenate with the template, which may be further used by other modules # x, event_x = x[:, :lens_x//2], x[:, lens_x//2:] # x = x[:, -lens_x//2:] aux_dict = { "attn": attn, "removed_indexes_s": removed_indexes_s, # used for visualization } return x, aux_dict def forward(self, z, x, event_z, event_x, ce_template_mask=None, ce_keep_rate=None, tnc_keep_rate=None, return_last_attn=False): x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,) return x, aux_dict def _create_vision_transformer(pretrained=False, **kwargs): model = VisionTransformerCE(**kwargs) if pretrained: if 'npz' in pretrained: model.load_pretrained(pretrained, prefix='') else: checkpoint = torch.load(pretrained, map_location="cpu") missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False) print('Load pretrained model from: ' + pretrained) return model def vit_base_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model def vit_large_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model 的train loss差不多0.4，val loss差不多1.5 。修改模型① # 将 4输入分开，构建新的相同模态结合的2输入，2分支 import math import logging from functools import partial from collections import OrderedDict from copy import deepcopy import torch import torch.nn as nn import torch.nn.functional as F from timm.models.layers import to_2tuple from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise from .utils import combine_tokens, recover_tokens from .vit import VisionTransformer from ..layers.attn_blocks import CEBlock _logger = logging.getLogger(__name__) class VisionTransformerCE(VisionTransformer): """ Vision Transformer with candidate elimination (CE) module A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 Includes distillation token & head support for `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877 """ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False, drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None, weight_init='', ce_loc=None, ce_keep_ratio=None): super().__init__() if isinstance(img_size, tuple): self.img_size = img_size else: self.img_size = to_2tuple(img_size) self.patch_size = patch_size self.in_chans = in_chans self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models self.num_tokens = 2 if distilled else 1 norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule blocks = [] ce_index = 0 self.ce_loc = ce_loc for i in range(depth): ce_keep_ratio_i = 1.0 if ce_loc is not None and i in ce_loc: ce_keep_ratio_i = ce_keep_ratio[ce_index] ce_index += 1 blocks.append( CEBlock( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, keep_ratio_search=ce_keep_ratio_i) ) self.blocks = nn.Sequential(*blocks) self.norm = norm_layer(embed_dim) self.init_weights(weight_init) def forward_features(self, z, x, event_z, event_x, mask_z=None, mask_x=None, ce_template_mask=None, ce_keep_rate=None, return_last_attn=False ): # 分支1 处理流程 B, H, W = x.shape[0], x.shape[2], x.shape[3] x = self.patch_embed(x) z = self.patch_embed(z) z += self.pos_embed_z x += self.pos_embed_x if mask_z is not None and mask_x is not None: mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_z = mask_z.flatten(1).unsqueeze(-1) mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_x = mask_x.flatten(1).unsqueeze(-1) mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode) mask_x = mask_x.squeeze(-1) if self.add_cls_token: cls_tokens = self.cls_token.expand(B, -1, -1) cls_tokens = cls_tokens + self.cls_pos_embed if self.add_sep_seg: x += self.search_segment_pos_embed z += self.template_segment_pos_embed x = combine_tokens(z, x, mode=self.cat_mode) if self.add_cls_token: x = torch.cat([cls_tokens, x], dim=1) x = self.pos_drop(x) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device) global_index_t = global_index_t.repeat(B, 1) global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device) global_index_s = global_index_s.repeat(B, 1) removed_indexes_s = [] for i, blk in enumerate(self.blocks): x, global_index_t, global_index_s, removed_index_s, attn = blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate) if self.ce_loc is not None and i in self.ce_loc: removed_indexes_s.append(removed_index_s) x = self.norm(x) # # [bs, n_patch, dim] = [bs, 320, 768] 320 = 64 + 256 # # 分支2 处理流程 event_x = self.pos_embed_event(event_x) event_z = self.pos_embed_event(event_z) event_x += self.pos_embed_x event_z += self.pos_embed_z event_x = combine_tokens(event_z, event_x, mode=self.cat_mode) if self.add_cls_token: event_x = torch.cat([cls_tokens, event_x], dim=1) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t1 = torch.linspace(0, lens_z - 1, lens_z).to(event_x.device) global_index_t1 = global_index_t1.repeat(B, 1) global_index_s1 = torch.linspace(0, lens_x - 1, lens_x).to(event_x.device) global_index_s1 = global_index_s1.repeat(B, 1) removed_indexes_s1 = [] for i, blk in enumerate(self.blocks): event_x, global_index_t1, global_index_s1, removed_index_s1, attn = blk(event_x, global_index_t1, global_index_s1, mask_x, ce_template_mask, ce_keep_rate) if self.ce_loc is not None and i in self.ce_loc: removed_indexes_s1.append(removed_index_s1) event_x = self.norm(event_x) # print('x.shape: ',x.shape) #x.shape: torch.Size([2, 320, 768]) # print('event_x.shape: ',event_x.shape) # event_x.shape: torch.Size([2, 320, 768]) x_cat = torch.cat([x, event_x], dim=1) # print('x_cat.shape: ',x_cat.shape) # x_cat.shape: torch.Size([2, 640, 768]) x = x_cat aux_dict = { "attn": attn, "removed_indexes_s": removed_indexes_s, # used for visualization } return x, aux_dict def forward(self, z, x, event_z, event_x, ce_template_mask=None, ce_keep_rate=None, tnc_keep_rate=None, return_last_attn=False): x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,) return x, aux_dict def _create_vision_transformer(pretrained=False, **kwargs): model = VisionTransformerCE(**kwargs) if pretrained: if 'npz' in pretrained: model.load_pretrained(pretrained, prefix='') else: checkpoint = torch.load(pretrained, map_location="cpu") missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False) print('Load pretrained model from: ' + pretrained) return model def vit_base_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model def vit_large_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model中train loss 0.5左右，val loss 1.2左右，代码② # 将 4输入分开，构建新的相同模态结合的2输入，2分支 import math import logging from functools import partial from collections import OrderedDict from copy import deepcopy import torch import torch.nn as nn import torch.nn.functional as F from timm.models.layers import to_2tuple from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise from .utils import combine_tokens, recover_tokens from .vit import VisionTransformer from ..layers.attn_blocks import CEBlock from .ad_counter_guide import Counter_Guide_Enhanced _logger = logging.getLogger(__name__) class VisionTransformerCE(VisionTransformer): """ Vision Transformer with candidate elimination (CE) module A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 Includes distillation token & head support for `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877 """ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False, drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None, weight_init='', ce_loc=None, ce_keep_ratio=None): super().__init__() if isinstance(img_size, tuple): self.img_size = img_size else: self.img_size = to_2tuple(img_size) self.patch_size = patch_size self.in_chans = in_chans self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models self.num_tokens = 2 if distilled else 1 norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule blocks = [] ce_index = 0 self.ce_loc = ce_loc for i in range(depth): ce_keep_ratio_i = 1.0 if ce_loc is not None and i in ce_loc: ce_keep_ratio_i = ce_keep_ratio[ce_index] ce_index += 1 blocks.append( CEBlock( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, keep_ratio_search=ce_keep_ratio_i) ) self.blocks = nn.Sequential(*blocks) self.norm = norm_layer(embed_dim) self.init_weights(weight_init) # 添加交互模块counter_guide # self.counter_guide = Counter_Guide(768, 768) self.counter_guide = Counter_Guide_Enhanced(768, 768) def forward_features(self, z, x, event_z, event_x, mask_z=None, mask_x=None, ce_template_mask=None, ce_keep_rate=None, return_last_attn=False ): # 分支1 处理流程 B, H, W = x.shape[0], x.shape[2], x.shape[3] x = self.patch_embed(x) z = self.patch_embed(z) z += self.pos_embed_z x += self.pos_embed_x if mask_z is not None and mask_x is not None: mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_z = mask_z.flatten(1).unsqueeze(-1) mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_x = mask_x.flatten(1).unsqueeze(-1) mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode) mask_x = mask_x.squeeze(-1) if self.add_cls_token: cls_tokens = self.cls_token.expand(B, -1, -1) cls_tokens = cls_tokens + self.cls_pos_embed if self.add_sep_seg: x += self.search_segment_pos_embed z += self.template_segment_pos_embed x = combine_tokens(z, x, mode=self.cat_mode) if self.add_cls_token: x = torch.cat([cls_tokens, x], dim=1) x = self.pos_drop(x) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device) global_index_t = global_index_t.repeat(B, 1) global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device) global_index_s = global_index_s.repeat(B, 1) removed_indexes_s = [] # # 分支2 处理流程 event_x = self.pos_embed_event(event_x) event_z = self.pos_embed_event(event_z) event_x += self.pos_embed_x event_z += self.pos_embed_z event_x = combine_tokens(event_z, event_x, mode=self.cat_mode) if self.add_cls_token: event_x = torch.cat([cls_tokens, event_x], dim=1) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t1 = torch.linspace(0, lens_z - 1, lens_z).to(event_x.device) global_index_t1 = global_index_t1.repeat(B, 1) global_index_s1 = torch.linspace(0, lens_x - 1, lens_x).to(event_x.device) global_index_s1 = global_index_s1.repeat(B, 1) removed_indexes_s1 = [] for i, blk in enumerate(self.blocks): # 第一个分支处理 x, global_index_t, global_index_s, removed_index_s, attn = blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate) # 第二个分支处理 event_x, global_index_t1, global_index_s1, removed_index_s1, attn1 = blk(event_x, global_index_t1, global_index_s1, mask_x, ce_template_mask, ce_keep_rate) if self.ce_loc is not None and i in self.ce_loc: removed_indexes_s.append(removed_index_s) removed_indexes_s1.append(removed_index_s1) # 在第1层和第2层增加counter_guide模块,验证早期融合效果 if i ==0 : enhanced_x, enhanced_event_x = self.counter_guide(x, event_x) # 将增强后的特征与原特征相加 x = x + enhanced_x event_x = event_x + enhanced_event_x # 应用LayerNorm归一化处理 x = self.norm(x) event_x = self.norm(event_x) x_cat = torch.cat([event_x,x], dim=1) x = x_cat aux_dict = { "attn": attn, "attn1": attn1, "removed_indexes_s": removed_indexes_s, # used for visualization "removed_indexes_s1": removed_indexes_s1, } return x, aux_dict def forward(self, z, x, event_z, event_x, ce_template_mask=None, ce_keep_rate=None, tnc_keep_rate=None, return_last_attn=False): x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,) return x, aux_dict def _create_vision_transformer(pretrained=False, **kwargs): model = VisionTransformerCE(**kwargs) if pretrained: if 'npz' in pretrained: model.load_pretrained(pretrained, prefix='') else: checkpoint = torch.load(pretrained, map_location="cpu") missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False) print('Load pretrained model from: ' + pretrained) return model def vit_base_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model def vit_large_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model 中train loss差不多1.27 ， val loss差不多1.45，

question:引入的counter_guide是：import torch import torch.nn as nn import torch.nn.functional as F # 为结合vit中的attention优势,设计了新的交叉注意力 class Cross_Attention(nn.Module): def __init__(self, input_channels, output_channels): super(Cross_Attention, self).__init__() self.query = nn.Linear(input_channels, output_channels) self.key = nn.Linear(input_channels, output_channels) self.value = nn.Linear(input_channels, output_channels) self.scale = output_channels ** -0.5 # 缩放因子，防止softmax的梯度消失/爆炸 def forward(self, assistant, present): # 生成查询、键和值 q = self.query(assistant) k = self.key(present) v = self.value(present) # 计算注意力分数 attention_scores = torch.bmm(q, k.transpose(1, 2)) * self.scale attention = F.softmax(attention_scores, dim=-1) # 应用注意力机制 attended_features = torch.bmm(attention, v) return attended_features class Multi_Context(nn.Module): def __init__(self, input_channels, output_channels): # 修正了def init为def init super(Multi_Context, self).__init__() self.linear1 = nn.Linear(input_channels, output_channels) self.linear2 = nn.Linear(input_channels, output_channels) self.linear3 = nn.Linear(input_channels, output_channels) self.linear_final = nn.Linear(output_channels * 3, output_channels) def forward(self, x): x1 = F.relu(self.linear1(x)) x2 = F.relu(self.linear2(x)) x3 = F.relu(self.linear3(x)) x = torch.cat([x1, x2, x3], dim=-1) # Concatenate along the feature dimension x = self.linear_final(x) return x class Adaptive_Weight(nn.Module): def __init__(self, input_channels): # 修正了def init为def init super(Adaptive_Weight, self).__init__() self.fc1 = nn.Linear(input_channels, input_channels // 4) self.fc2 = nn.Linear(input_channels // 4, input_channels) self.sigmoid = nn.Sigmoid() def forward(self, x): x_avg = torch.mean(x, dim=1) # Global Average Pooling along the ‘sequence length’ dimension weight = F.relu(self.fc1(x_avg)) weight = self.fc2(weight) weight = self.sigmoid(weight).unsqueeze(1) out = x * weight return out class Counter_Guide_Enhanced(nn.Module): def __init__(self, input_channels, output_channels): super(Counter_Guide_Enhanced, self).__init__() self.mc = Multi_Context(input_channels, output_channels) self.cross_attention = Cross_Attention(output_channels, output_channels) # 结合自适应 self.adaptive_weight = Adaptive_Weight(output_channels) def forward(self, x, event_x): # 首先使用Multi-Context增强事件特征 enhanced_event_x = self.mc(event_x) # 使用Cross-Attention进行特征交互,(x和event_x互相增强) enhanced_x = self.cross_attention(x, enhanced_event_x) enhanced_event_x_via_x = self.cross_attention(event_x, enhanced_x) # 进一步通过adaptive weight调整增强的特征权重 adaptive_enhanced_x = self.adaptive_weight(enhanced_x) adaptive_enhanced_event_x = self.adaptive_weight(enhanced_event_x_via_x) return adaptive_enhanced_x, adaptive_enhanced_event_x

question:但是引入该模块之后，按照原始的训练配置，train loss下降，但是下降至1.5左右，但是未加Counter_Guide_Enhanced之后，train loss下降可以下降至0.4左右

question:那么现在考虑简化Counter_Guide_Enhanced中的输出维度。首先代码① import torch import torch.nn as nn import torch.nn.functional as F # 为结合vit中的attention优势,设计了新的交叉注意力 class Cross_Attention(nn.Module): def __init__(self, input_channels, output_channels): super(Cross_Attention, self).__init__() self.query = nn.Linear(input_channels, output_channels) self.key = nn.Linear(input_channels, output_channels) self.value = nn.Linear(input_channels, output_channels) self.scale = output_channels ** -0.5 # 缩放因子，防止softmax的梯度消失/爆炸 def forward(self, assistant, present): # 生成查询、键和值 q = self.query(assistant) k = self.key(present) v = self.value(present) # 计算注意力分数 attention_scores = torch.bmm(q, k.transpose(1, 2)) * self.scale attention = F.softmax(attention_scores, dim=-1) # 应用注意力机制 attended_features = torch.bmm(attention, v) return attended_features class Multi_Context(nn.Module): def __init__(self, input_channels, output_channels): # 修正了def init为def init super(Multi_Context, self).__init__() self.linear1 = nn.Linear(input_channels, output_channels) self.linear2 = nn.Linear(input_channels, output_channels) self.linear3 = nn.Linear(input_channels, output_channels) self.linear_final = nn.Linear(output_channels * 3, output_channels) def forward(self, x): x1 = F.relu(self.linear1(x)) x2 = F.relu(self.linear2(x)) x3 = F.relu(self.linear3(x)) x = torch.cat([x1, x2, x3], dim=-1) # Concatenate along the feature dimension x = self.linear_final(x) return x class Adaptive_Weight(nn.Module): def __init__(self, input_channels): # 修正了def init为def init super(Adaptive_Weight, self).__init__() self.fc1 = nn.Linear(input_channels, input_channels // 4) self.fc2 = nn.Linear(input_channels // 4, input_channels) self.sigmoid = nn.Sigmoid() def forward(self, x): x_avg = torch.mean(x, dim=1) # Global Average Pooling along the ‘sequence length’ dimension weight = F.relu(self.fc1(x_avg)) weight = self.fc2(weight) weight = self.sigmoid(weight).unsqueeze(1) out = x * weight return out class Counter_Guide_Enhanced(nn.Module): def __init__(self, input_channels, output_channels): super(Counter_Guide_Enhanced, self).__init__() self.mc = Multi_Context(input_channels, output_channels) self.cross_attention = Cross_Attention(output_channels, output_channels) # 结合自适应 self.adaptive_weight = Adaptive_Weight(output_channels) def forward(self, x, event_x): # 首先使用Multi-Context增强事件特征 enhanced_event_x = self.mc(event_x) # 使用Cross-Attention进行特征交互,(x和event_x互相增强) enhanced_x = self.cross_attention(x, enhanced_event_x) enhanced_event_x_via_x = self.cross_attention(event_x, enhanced_x) # 进一步通过adaptive weight调整增强的特征权重 adaptive_enhanced_x = self.adaptive_weight(enhanced_x) adaptive_enhanced_event_x = self.adaptive_weight(enhanced_event_x_via_x) return adaptive_enhanced_x, adaptive_enhanced_event_x 调用代码是：# 将 4输入分开，构建新的相同模态结合的2输入，2分支 import math import logging from functools import partial from collections import OrderedDict from copy import deepcopy import torch import torch.nn as nn import torch.nn.functional as F from timm.models.layers import to_2tuple from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise from .utils import combine_tokens, recover_tokens from .vit import VisionTransformer from ..layers.attn_blocks import CEBlock from .new_counter_guide import Counter_Guide from .ad_counter_guide import Counter_Guide_Enhanced _logger = logging.getLogger(__name__) class VisionTransformerCE(VisionTransformer): """ Vision Transformer with candidate elimination (CE) module A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` - https://arxiv.org/abs/2010.11929 Includes distillation token & head support for `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877 """ def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False, drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None, weight_init='', ce_loc=None, ce_keep_ratio=None): super().__init__() if isinstance(img_size, tuple): self.img_size = img_size else: self.img_size = to_2tuple(img_size) self.patch_size = patch_size self.in_chans = in_chans self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models self.num_tokens = 2 if distilled else 1 norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU self.patch_embed = embed_layer( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4) dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule blocks = [] ce_index = 0 self.ce_loc = ce_loc for i in range(depth): ce_keep_ratio_i = 1.0 if ce_loc is not None and i in ce_loc: ce_keep_ratio_i = ce_keep_ratio[ce_index] ce_index += 1 blocks.append( CEBlock( dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, keep_ratio_search=ce_keep_ratio_i) ) self.blocks = nn.Sequential(*blocks) self.norm = norm_layer(embed_dim) self.init_weights(weight_init) # 添加交互模块counter_guide # self.counter_guide = Counter_Guide(768, 768) self.counter_guide = Counter_Guide_Enhanced(768, 768) def forward_features(self, z, x, event_z, event_x, mask_z=None, mask_x=None, ce_template_mask=None, ce_keep_rate=None, return_last_attn=False ): # 分支1 处理流程 B, H, W = x.shape[0], x.shape[2], x.shape[3] x = self.patch_embed(x) z = self.patch_embed(z) z += self.pos_embed_z x += self.pos_embed_x if mask_z is not None and mask_x is not None: mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_z = mask_z.flatten(1).unsqueeze(-1) mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0] mask_x = mask_x.flatten(1).unsqueeze(-1) mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode) mask_x = mask_x.squeeze(-1) if self.add_cls_token: cls_tokens = self.cls_token.expand(B, -1, -1) cls_tokens = cls_tokens + self.cls_pos_embed if self.add_sep_seg: x += self.search_segment_pos_embed z += self.template_segment_pos_embed x = combine_tokens(z, x, mode=self.cat_mode) if self.add_cls_token: x = torch.cat([cls_tokens, x], dim=1) x = self.pos_drop(x) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device) global_index_t = global_index_t.repeat(B, 1) global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device) global_index_s = global_index_s.repeat(B, 1) removed_indexes_s = [] # # 分支2 处理流程 event_x = self.pos_embed_event(event_x) event_z = self.pos_embed_event(event_z) event_x += self.pos_embed_x event_z += self.pos_embed_z event_x = combine_tokens(event_z, event_x, mode=self.cat_mode) if self.add_cls_token: event_x = torch.cat([cls_tokens, event_x], dim=1) lens_z = self.pos_embed_z.shape[1] lens_x = self.pos_embed_x.shape[1] global_index_t1 = torch.linspace(0, lens_z - 1, lens_z).to(event_x.device) global_index_t1 = global_index_t1.repeat(B, 1) global_index_s1 = torch.linspace(0, lens_x - 1, lens_x).to(event_x.device) global_index_s1 = global_index_s1.repeat(B, 1) removed_indexes_s1 = [] for i, blk in enumerate(self.blocks): # 第一个分支处理 x, global_index_t, global_index_s, removed_index_s, attn = blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate) # 第二个分支处理 event_x, global_index_t1, global_index_s1, removed_index_s1, attn = blk(event_x, global_index_t1, global_index_s1, mask_x, ce_template_mask, ce_keep_rate) if self.ce_loc is not None and i in self.ce_loc: removed_indexes_s.append(removed_index_s) removed_indexes_s1.append(removed_index_s1) # 在第1层和第2层增加counter_guide模块,验证早期融合效果 if i == 0 : enhanced_x, enhanced_event_x = self.counter_guide(x, event_x) # 将增强后的特征与原特征相加 x = x + enhanced_x event_x = event_x + enhanced_event_x # 应用LayerNorm归一化处理 x = self.norm(x) event_x = self.norm(event_x) x_cat = torch.cat([event_x,x], dim=1) x = x_cat aux_dict = { "attn": attn, "removed_indexes_s": removed_indexes_s, # used for visualization } return x, aux_dict def forward(self, z, x, event_z, event_x, ce_template_mask=None, ce_keep_rate=None, tnc_keep_rate=None, return_last_attn=False): x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate,) return x, aux_dict def _create_vision_transformer(pretrained=False, **kwargs): model = VisionTransformerCE(**kwargs) if pretrained: if 'npz' in pretrained: model.load_pretrained(pretrained, prefix='') else: checkpoint = torch.load(pretrained, map_location="cpu") missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=False) print('Load pretrained model from: ' + pretrained) return model def vit_base_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model def vit_large_patch16_224_ce(pretrained=False, **kwargs): """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929). """ model_kwargs = dict( patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs) model = _create_vision_transformer(pretrained=pretrained, **model_kwargs) return model ，那么按照要求，如何降低输出维度？