question:上面的代码中loss_rank处理的对象是什么？

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 import random import numpy as np _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 masking_fea(self,z, event_z, x, event_x, ratio=0.8 ): b,nz,c = z.shape b,nez,c = event_z.shape b,nx,c = x.shape b,nex,c = event_x.shape assert(nz == nez) assert(nx == nex) lenz_out = int(nz*ratio) lenx_out = int(nx*ratio) mask_nz = torch.rand(b,nz).float() mask_ez = torch.rand(b,nez).float() mask_nx = torch.rand(b,nx).float() mask_ex = torch.rand(b,nex).float() mask_nz = mask_nz>0.4 mask_ez = mask_ez>0.4 mask_ez = ~mask_nz + mask_ez mask_nz_idx = mask_nz.float().sort(1,descending=True)[-1].to(device = z.device) mask_ez_idx = mask_ez.float().sort(1,descending=True)[-1].to(device = z.device) mask_nx = mask_nx>0.4 mask_ex = mask_ex>0.4 mask_ex = ~mask_nx + mask_ex mask_nx_idx = mask_nx.float().sort(1,descending=True)[-1].to(device = z.device) mask_ex_idx = mask_ex.float().sort(1,descending=True)[-1].to(device = z.device) masked_z = torch.gather(z, 1, mask_nz_idx[:,:lenz_out,None].repeat([1,1,c])) masked_ez = torch.gather(event_z, 1, mask_ez_idx[:,:lenz_out,None].repeat([1,1,c])) masked_x = torch.gather(x, 1, mask_nx_idx[:,:lenx_out,None].repeat([1,1,c])) masked_ex = torch.gather(event_x, 1, mask_ex_idx[:,:lenx_out,None].repeat([1,1,c])) return masked_z, masked_ez, masked_x, masked_ex,{'x1':mask_nx_idx[:,:lenx_out],'x0':mask_nx_idx[:,lenx_out:], 'ex1':mask_ex_idx[:,:lenx_out],'ex0':mask_ex_idx[:,lenx_out:], } 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,Track=False ): B, H, W = x.shape[0], x.shape[2], x.shape[3] # print('shape of event_z before projection:{}, event_x:{}'.format(event_z.shape, event_x.shape)) 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) # print('shape of event_z:{}, event_x:{}, x:{}, z:{}'.format(event_z.shape,event_x.shape,x.shape,z.shape )) 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 if Track == False: z, event_z, x, event_x, token_idx = self.masking_fea(z, event_z, x, event_x, ratio=0.9) 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] lens_z = z.shape[1] lens_x = 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 = [] out_attn = [] for i, blk in enumerate(self.blocks): # out_global_s.append(global_index_s) # out_global_t.append(global_index_t) 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) out_attn.append(attn) # print('shape of attn:{}, lens_z:{}, lens_x:{}'.format(attn.shape, lens_z, lens_x)) out_attn_idx = random.choice(np.arange(len(out_attn))) out_attn = out_attn[out_attn_idx] 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 Track == False: idx1 = token_idx['x1'] idx0 = token_idx['x0'] idex1 = token_idx['ex1'] idex0 = token_idx['ex0'] ex = x[:,idex1.shape[1]:] x = x[:,:idex1.shape[1]] # 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 = idx0.shape[1] 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([idx1, idx0], 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) ex = torch.cat([ex, pad_x], dim=1) index_all = torch.cat([idex1, idex0], dim=1) # recover original token order C = ex.shape[-1] ex = torch.zeros_like(ex).scatter_(dim=1, index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64), src=ex) x = torch.cat([x,ex],dim=1) x = recover_tokens(x, lens_z_new, lens_x, mode=self.cat_mode) event_x = x[:, lens_x:] # RGB head x = x[:, :lens_x] # RGB head x = torch.cat([event_x, x], dim=1) aux_dict = { # "attn": attn, "attn": out_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,Track=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,Track=Track) 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 和 ④ # 将 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 from .ad_counter_guide_downdim 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_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

question:那么综上，代码①对应于代码④ ，代码②对应于代码③ ，那么描述一下对饮关系

question:代码③对应的是：from . import BaseActor from lib.utils.misc import NestedTensor from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy import torch from lib.utils.merge import merge_template_search from ...utils.heapmap_utils import generate_heatmap from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate class CEUTrackActor(BaseActor): """ Actor for training CEUTrack models """ def __init__(self, net, objective, loss_weight, settings, cfg=None): super().__init__(net, objective) self.loss_weight = loss_weight self.settings = settings self.bs = self.settings.batchsize # batch size self.cfg = cfg def __call__(self, data): """ args: data - The input data, should contain the fields 'template', 'search', 'gt_bbox'. template_images: (N_t, batch, 3, H, W) search_images: (N_s, batch, 3, H, W) returns: loss - the training loss status - dict containing detailed losses """ # forward pass out_dict = self.forward_pass(data) # compute losses loss, status = self.compute_losses(out_dict, data) return loss, status def forward_pass(self, data): # currently only support 1 template and 1 search region assert len(data['template_images']) == 1 assert len(data['search_images']) == 1 assert len(data['template_event']) == 1 assert len(data['search_event']) == 1 template_list = [] for i in range(self.settings.num_template): template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:]) # (batch, 3, 128, 128) # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:]) # (batch, 128, 128) template_list.append(template_img_i) search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:]) # (batch, 3, 320, 320) # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:]) # (batch, 320, 320) template_event = data['template_event'][0].view(-1, *data['template_event'].shape[2:]) search_event = data['search_event'][0].view(-1, *data['search_event'].shape[2:]) box_mask_z = None ce_keep_rate = None if self.cfg.MODEL.BACKBONE.CE_LOC: box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device, data['template_anno'][0]) ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch, total_epochs=ce_start_epoch + ce_warm_epoch, ITERS_PER_EPOCH=1, base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0]) if len(template_list) == 1: template_list = template_list[0] out_dict = self.net(template=template_list, search=search_img, event_template=template_event, event_search=search_event, ce_template_mask=box_mask_z, ce_keep_rate=ce_keep_rate, return_last_attn=False) return out_dict def compute_losses(self, pred_dict, gt_dict, return_status=True): # gt gaussian map gt_bbox = gt_dict['search_anno'][-1] # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4) gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE) gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1) # Get boxes pred_boxes = pred_dict['pred_boxes'] if torch.isnan(pred_boxes).any(): raise ValueError("Network outputs is NAN! Stop Training") num_queries = pred_boxes.size(1) pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4) # (B,N,4) --> (BN,4) (x1,y1,x2,y2) gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0) # (B,4) --> (B,1,4) --> (B,N,4) # compute giou and iou try: giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec) # (BN,4) (BN,4) except: giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda() # compute l1 loss l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec) # (BN,4) (BN,4) # compute location loss if 'score_map' in pred_dict: location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps) else: location_loss = torch.tensor(0.0, device=l1_loss.device) # weighted sum loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss if return_status: # status for log mean_iou = iou.detach().mean() status = {"Loss/total": loss.item(), "Loss/giou": giou_loss.item(), "Loss/l1": l1_loss.item(), "Loss/location": location_loss.item(), "IoU": mean_iou.item()} return loss, status else: return loss 代码④ 对应的是 from . import BaseActor from lib.utils.misc import NestedTensor from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy import torch from lib.utils.merge import merge_template_search from ...utils.heapmap_utils import generate_heatmap from ...utils.ce_utils import generate_mask_cond, adjust_keep_rate class CEUTrackActor(BaseActor): """ Actor for training CEUTrack models """ def __init__(self, net, objective, loss_weight, settings, cfg=None): super().__init__(net, objective) self.loss_weight = loss_weight self.settings = settings self.bs = self.settings.batchsize # batch size self.cfg = cfg def __call__(self, data): """ args: data - The input data, should contain the fields 'template', 'search', 'gt_bbox'. template_images: (N_t, batch, 3, H, W) search_images: (N_s, batch, 3, H, W) returns: loss - the training loss status - dict containing detailed losses """ # forward pass out_dict = self.forward_pass(data) # compute losses loss, status = self.compute_losses(out_dict, data) return loss, status def forward_pass(self, data): # currently only support 1 template and 1 search region assert len(data['template_images']) == 1 assert len(data['search_images']) == 1 assert len(data['template_event']) == 1 assert len(data['search_event']) == 1 template_list = [] for i in range(self.settings.num_template): template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:]) # (batch, 3, 128, 128) # template_att_i = data['template_att'][i].view(-1, *data['template_att'].shape[2:]) # (batch, 128, 128) template_list.append(template_img_i) search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:]) # (batch, 3, 320, 320) # search_att = data['search_att'][0].view(-1, *data['search_att'].shape[2:]) # (batch, 320, 320) template_event = data['template_event'][0].view(-1, *data['template_event'].shape[2:]) search_event = data['search_event'][0].view(-1, *data['search_event'].shape[2:]) box_mask_z = None ce_keep_rate = None if self.cfg.MODEL.BACKBONE.CE_LOC: box_mask_z = generate_mask_cond(self.cfg, template_list[0].shape[0], template_list[0].device, data['template_anno'][0]) ce_start_epoch = self.cfg.TRAIN.CE_START_EPOCH ce_warm_epoch = self.cfg.TRAIN.CE_WARM_EPOCH ce_keep_rate = adjust_keep_rate(data['epoch'], warmup_epochs=ce_start_epoch, total_epochs=ce_start_epoch + ce_warm_epoch, ITERS_PER_EPOCH=1, base_keep_rate=self.cfg.MODEL.BACKBONE.CE_KEEP_RATIO[0]) if len(template_list) == 1: template_list = template_list[0] out_dict = self.net(template=template_list, search=search_img, event_template=template_event, event_search=search_event, ce_template_mask=box_mask_z, ce_keep_rate=ce_keep_rate, return_last_attn=False) return out_dict def compute_losses(self, pred_dict, gt_dict, return_status=True): # gt gaussian map gt_bbox = gt_dict['search_anno'][-1] # (Ns, batch, 4) (x1,y1,w,h) -> (batch, 4) gt_gaussian_maps = generate_heatmap(gt_dict['search_anno'], self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE) gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1) # Get boxes pred_boxes = pred_dict['pred_boxes'] if torch.isnan(pred_boxes).any(): raise ValueError("Network outputs is NAN! Stop Training") num_queries = pred_boxes.size(1) pred_boxes_vec = box_cxcywh_to_xyxy(pred_boxes).view(-1, 4) # (B,N,4) --> (BN,4) (x1,y1,x2,y2) gt_boxes_vec = box_xywh_to_xyxy(gt_bbox)[:, None, :].repeat((1, num_queries, 1)).view(-1, 4).clamp(min=0.0, max=1.0) # (B,4) --> (B,1,4) --> (B,N,4) # compute giou and iou try: giou_loss, iou = self.objective['giou'](pred_boxes_vec, gt_boxes_vec) # (BN,4) (BN,4) except: giou_loss, iou = torch.tensor(0.0).cuda(), torch.tensor(0.0).cuda() # compute l1 loss l1_loss = self.objective['l1'](pred_boxes_vec, gt_boxes_vec) # (BN,4) (BN,4) # compute location loss if 'score_map' in pred_dict: location_loss = self.objective['focal'](pred_dict['score_map'], gt_gaussian_maps) else: location_loss = torch.tensor(0.0, device=l1_loss.device) rank_loss = self.loss_rank(pred_dict,gt_dict['search_anno'], gt_dict['template_anno']) # weighted sum loss = self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss + rank_loss*1.2 if return_status: # status for log mean_iou = iou.detach().mean() status = {"Loss/total": loss.item(), "Loss/giou": giou_loss.item(), "Loss/l1": l1_loss.item(), "Loss/location": location_loss.item(), "IoU": mean_iou.item()} return loss, status else: return loss def _random_permute(self,matrix): # matrix = random.choice(matrix) b, c, h, w = matrix.shape idx = [ torch.randperm(c).to(matrix.device) for i in range(b)] idx = torch.stack(idx, dim=0)[:, :, None, None].repeat([1,1,h,w]) # idx = torch.randperm(c)[None,:,None,None].repeat([b,1,h,w]).to(matrix.device) matrix01 = torch.gather(matrix, 1, idx) return matrix01 def crop_flag(self, flag, global_index_s, global_index_t,H1 = 64, H2 = 256): B,Ls = global_index_s.shape B, Lt = global_index_t.shape B,C,L1,L2 = flag.shape flag_t = flag[:,:,:H1,:] flag_s = flag[:,:,H1:,:] flag_t = torch.gather(flag_t,2,global_index_t[:,None,:,None].repeat([1,C,1,L2]).long()) flag_s = torch.gather(flag_s,2,global_index_s[:,None,:,None].repeat([1,C,1,L2]).long()) flag = torch.cat([flag_t, flag_s], dim = 2) flag_t = flag[:,:,:,:H1] flag_s = flag[:,:,:,H1:] flag_t = torch.gather(flag_t,3,global_index_t[:,None,None,:].repeat([1,C,int(Ls+Lt),1]).long()) flag_s = torch.gather(flag_s,3,global_index_s[:,None,None,:].repeat([1,C,int(Ls+Lt),1]).long()) flag = torch.cat([flag_t, flag_s], dim = 3) B, C, L11, L12 = flag.shape try: assert(L11 == int(Lt + Ls)) assert(L12 == int(Lt + Ls)) except: print('L11:{}, L12:{}, L1:{}, L2:{}'.format(L11, L12, L1, L2)) return flag def crop_fusion(self, flag, attn, global_index_s, global_index_t,H1 = 64, H2 = 256 ): flag = self.crop_flag(flag=flag, global_index_s=global_index_s, global_index_t=global_index_t) B,C,L1,L2 = flag.shape Ba, Ca, La, La2 = attn.shape _,idx1 = flag.mean(dim=3,keepdim=False).sort(dim=2,descending=True) # print('shape of flag:{}, idx1:{}'.format(flag.shape, idx1[:,:,:32,None].repeat([1,Ca,1,L2]).shape)) flag = torch.gather(flag,2,idx1[:,:,:32,None].repeat([1,C,1,L2]).long()) attn = torch.gather(attn,2,idx1[:,:,:32,None].repeat([1,Ca,1,L2]).long()) _,idx2 = flag.mean(dim=2,keepdim=False).sort(dim=2,descending=True) flag = torch.gather(flag,3,idx2[:,:,None,:32].repeat([1,C,32,1]).long()) attn = torch.gather(attn,3,idx2[:,:,None,:32].repeat([1,Ca,32,1]).long()) return attn * flag def loss_rank(self, outputs, targetsi, temp_annoi=None): """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4] The target boxes are expected in format (center_x, center_y, h, w), normalized by the image size. """ attn = outputs['attn'] # print('attn shape:{}'.format(attn.shape)) attn1 = torch.cat([attn[:,:,114:344,57:114], attn[:,:,114:344,344:]],dim=3) attn1 = attn1.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True) attn2 = torch.cat([attn[:,:,344:,:57], attn[:,:,344:,114:344]],dim=3) attn2 = attn2.mean(dim=0, keepdim=True).mean(dim=1, keepdim=True) # print('attn1 shape:{},attn2 shape:{}, attn:{}'.format(attn1.shape,attn2.shape,attn.shape)) # attn = self._random_permute(attn) # attn = attn[:,:,:,:] # B1, C1, H1, W1 = attn.shape # global_index_s = outputs['out_global_s'] # global_index_t = outputs['out_global_t'] # try: # assert((global_index_s.shape[1] + global_index_t.shape[1])== int(H1/2)) # except: # print('Falut,shape of attn:{}, s:{}, t:{}'.format(attn.shape,global_index_s.shape, global_index_t.shape )) # H1 = int(64) # H2 = int(256) # l_t = int(math.sqrt(64)) # l_s = int(math.sqrt(256)) # temp_anno = temp_annoi[0,:,:] # targets = targetsi[0,:,:] # r_s = torch.arange(l_s).to(temp_anno.device) # r_t = torch.arange(l_t).to(temp_anno.device) # r_t = r_t[None,:].repeat([B1,1]) # cx, cy, w, h = temp_anno[:,0:1], temp_anno[:,1:2], temp_anno[:,2:3], temp_anno[:,3:4] # cx *= l_t # cy *= l_t # w *= l_t # h *= l_t # flagx_01 = r_t >= cx - w/2 # flagx_02 = r_t <= cx + w/2 # flagy_02 = r_t >= cy - h/2 # flagy_01 = r_t <= cy + h/2 # flagx = flagx_01.float()*flagx_02.float() # flagy = flagy_01.float()*flagy_02.float() # flagx = flagx[:,None,:].repeat([1,l_t,1]) # flagy = flagy[:,:,None].repeat([1,1,l_t]) # flag = flagx*flagy # flagt = flag.reshape([B1, H1]) # cx, cy, w, h = targets[:,0:1], targets[:,1:2], targets[:,2:3], targets[:,3:4] # cx *= l_s # cy *= l_s # w *= l_s # h *= l_s # flagx_01 = r_s >= cx - w/2 # flagx_02 = r_s <= cx + w/2 # flagy_02 = r_s >= cy - h/2 # flagy_01 = r_s <= cy + h/2 # flagx = flagx_01.float()*flagx_02.float() # flagy = flagy_01.float()*flagy_02.float() # flagx = flagx[:,None,:].repeat([1,l_s,1]) # flagy = flagy[:,:,None].repeat([1,1,l_s]) # flag = flagx*flagy # flags = flag.reshape([B1, H2]) # flag = torch.cat([flagt, flags], dim=1) # flag_total = flag[:,:,None].repeat([1,1,int(H1+H2)]) * flag[:,None,:].repeat([1,int(H1+H2),1]) # attn1 = self.crop_fusion(flag_total[:,None,:,:], attn, global_index_s, global_index_t) attn = torch.cat([attn1, attn2],dim=1) B, C, H, W = attn.shape # _,s1,_ = torch.svd(attn1.reshape([B*C, H, W])) _,s1,_ = torch.svd(attn.reshape([B*C, H, W])) s01 = torch.abs(s1 - 1) return torch.mean(s01)，那么分析文件之间的对应关系