from . import BaseActor from lib.utils.box_ops import box_cxcywh_to_xyxy, box_xywh_to_xyxy, box_xyxy_to_cxcywh, box_cxcywh_to_xyxy, box_iou import torch from lib.utils.heapmap_utils import generate_heatmap from lib.utils.ce_utils import generate_mask_cond, adjust_keep_rate,generate_bbox_mask from lib.train.admin import multigpu import torch.nn as nn from lib.utils.misc import NestedTensor class ATCTrackActor(BaseActor): """ Actor for training the atctrack""" def __init__(self, net, objective, loss_weight, settings, cfg): super().__init__(net, objective) self.loss_weight = loss_weight self.settings = settings self.bs = self.settings.batchsize # batch size self.cfg = cfg self.task_cls_loss_fn = nn.CrossEntropyLoss() # reg loss self.confidence_reg_loss = nn.MSELoss() # sub mask index pred loss self.sub_mask_index_cls_loss = nn.BCELoss() def fix_bns(self): net = self.net.module if multigpu.is_multi_gpu(self.net) else self.net net.box_head.apply(self.fix_bn) def fix_bn(self, m): classname = m.__class__.__name__ if classname.find('BatchNorm') != -1: m.eval() def __call__(self, data): """ args: data - The input data, should contain the fields 'template', 'search', 'search_anno'. 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): # assert len(data['template_images']) == 1 template_list, search_list = [], [] for i in range(self.settings.num_template): template_img_i = data['template_images'][i].view(-1, *data['template_images'].shape[2:]) # (batch, 6, 128, 128) template_list.append(template_img_i) # search_img = data['search_images'][0].view(-1, *data['search_images'].shape[2:]) # (batch, 6, 320, 320) for i in range(self.settings.num_search): search_img_i = data['search_images'][i].view(-1, *data['search_images'].shape[2:]) search_list.append(search_img_i) # soft token type infor bbox_mask_list = [] for template_item in data["template_anno"]: template_bbox = template_item * template_list[0].shape[2] bbox_mask = torch.zeros((template_list[0].shape[0], template_list[0].shape[2], template_list[0].shape[3] )).to(template_list[0].device) bbox_mask = generate_bbox_mask(bbox_mask, template_bbox ) bbox_mask = bbox_mask.unfold(1, 16, 16).unfold(2, 16, 16) bbox_mask = bbox_mask.mean(dim=(-1, -2)).view(bbox_mask.shape[0],-1).unsqueeze(-1) bbox_mask_list.append(bbox_mask) ## nlp + subject mask exp_str_subject_mask_infor = data["nlp"] exp_str_list = [] subject_mask_list = [] for item in exp_str_subject_mask_infor: item_list = item.split("+") exp_str_list.append(item_list[0]) index_list = list(map(int, item_list[-1].split(","))) subject_mask_list.append(index_list) target_state_template_bbox = data["template_anno"][-2].view(-1, 4) target_state_new_template_bbox = data["template_anno"][-1].view(-1, 4) # ---- two-stage teacher labeling: expose frame identity for cache lookup ---- seq_names = data.get("seq_name", None) if isinstance(seq_names, str): seq_names = [seq_names] template_frame_ids_raw = data.get("template_frame_ids", None) target_state_seq_names = seq_names target_state_template_frame_ids = ( template_frame_ids_raw.T if template_frame_ids_raw is not None else None ) # (num_template, B) -> (B, num_template) out_dict = self.net(template=template_list, search=search_list, soft_token_template_mask = bbox_mask_list, exp_str=exp_str_list, exp_subject_mask = subject_mask_list, target_state_template_bbox=target_state_template_bbox, target_state_new_template_bbox=target_state_new_template_bbox, target_state_object_name=data.get("test_class", None), target_state_seq_names=target_state_seq_names, target_state_template_frame_ids=target_state_template_frame_ids, ) 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_bbox = gt_dict['search_anno'].view(-1, 4) gts = gt_bbox.unsqueeze(0) gt_gaussian_maps = generate_heatmap(gts, self.cfg.DATA.SEARCH.SIZE, self.cfg.MODEL.BACKBONE.STRIDE) gt_gaussian_maps = gt_gaussian_maps[-1].unsqueeze(1) # (B,1,H,W) # 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) ## involve confidence_pred_score confidence_pred = pred_dict["confidence_pred"].squeeze(1) confidence_loss = self.confidence_reg_loss(confidence_pred.float(), iou.float()) index_cls_loss = self.sub_mask_index_cls_loss(pred_dict["subject_infor_mask_pred"].squeeze(-1), pred_dict["subject_infor_mask_gt"]) # weighted sum if getattr(self.cfg.TRAIN, "TYPE", None) == "target_state": index_cls_weight = 0.0 else: index_cls_weight = 0.2 qwen_format_loss = pred_dict.get("qwen_format_loss", None) if qwen_format_loss is None: qwen_format_loss = torch.tensor(0.0, device=l1_loss.device) qwen_teacher_loss = pred_dict.get("qwen_teacher_loss", None) if qwen_teacher_loss is None: qwen_teacher_loss = torch.tensor(0.0, device=l1_loss.device) qwen_format_weight = getattr(self.cfg.TRAIN, "QWEN_FORMAT_WEIGHT", 0.0) qwen_teacher_weight = getattr(self.cfg.TRAIN, "QWEN_TEACHER_WEIGHT", 0.0) qwen_teacher_acc = None qwen_decision = pred_dict.get("target_state_update_decision", None) qwen_teacher_labels = pred_dict.get("qwen_teacher_labels", None) if qwen_decision is not None and qwen_teacher_labels is not None: valid_teacher = qwen_teacher_labels >= 0 if valid_teacher.any(): qwen_pred_labels = qwen_decision.to(device=qwen_teacher_labels.device, dtype=torch.long) qwen_teacher_acc = (qwen_pred_labels[valid_teacher] == qwen_teacher_labels[valid_teacher]).float().mean() loss = (self.loss_weight['giou'] * giou_loss + self.loss_weight['l1'] * l1_loss + self.loss_weight['focal'] * location_loss + confidence_loss + index_cls_loss * index_cls_weight + qwen_format_loss * qwen_format_weight + qwen_teacher_loss * qwen_teacher_weight) # 计算 index_cls 的准确率 predicted = pred_dict["subject_infor_mask_pred"].squeeze(-1) > 0.5 # 使用阈值0.5将logits转换为0或1 num = pred_dict["subject_infor_mask_gt"].numel() index_cls_acc = (predicted == pred_dict["subject_infor_mask_gt"]).sum().item() / num 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/confidence_loss": confidence_loss.item(), "Loss/location": location_loss.item(), "index_cls_loss": index_cls_loss.item(), "Loss/qwen_format": qwen_format_loss.item(), "Loss/qwen_teacher": qwen_teacher_loss.item(), "Loss/qwen_format_weighted": (qwen_format_loss * qwen_format_weight).item(), "Loss/qwen_teacher_weighted": (qwen_teacher_loss * qwen_teacher_weight).item(), "Qwen/teacher_acc": qwen_teacher_acc.item() if qwen_teacher_acc is not None else -1.0, "index_cls_acc": index_cls_acc, "IoU_main": mean_iou.item() } return loss, status else: return loss