import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from pytorch_metric_learning import distances def tr_loss(batch_dict,key): loss1 = (batch_dict[key][:, 0:3, 3] - batch_dict['pose_to_frame'][:, 0:3, 3]).norm(dim=1).mean() loss2 = (torch.acos(torch.clip(batch_dict[key][:, 0, 0].view(-1, 1), -1, 1)) - torch.acos(torch.clip(batch_dict['pose_to_frame'][:, 0, 0].view(-1, 1), -1, 1))).norm(dim=1).mean() / 3.1415 * 180 return loss1, loss2 def gen_points_loss(batch_dict): key_points_gen = batch_dict['key_points_gen'] key_points = batch_dict['key_points'] key_points_gen1 = torch.cat((key_points_gen, key_points_gen * 0), dim=2) key_points_gen1[:, :, 3] = 1 # pose_query=batch_dict['pose_query'] # pose_positive=batch_dict['pose_positive'] # poses=torch.cat((pose_query,pose_positive),dim=0) # key_points_gen2=torch.bmm(poses,key_points_gen1.permute(0,2,1)).permute(0,2,1) # key_points2=torch.bmm(poses,key_points.permute(0,2,1)).permute(0,2,1) # loss_gpo=(key_points_gen2[:,:,:2]-key_points2[:,:,:2]).norm(p=1,dim=2).mean() pose_to_frame = batch_dict['pose_to_frame'] B = pose_to_frame.shape[0] src_pts = key_points[:B] tgt_pts = key_points[B:] src_pts_gen = key_points_gen1[:B] tgt_pts_gen = key_points_gen1[B:] srcs = torch.cat((src_pts, src_pts_gen), dim=0) tgts = torch.cat((tgt_pts_gen, tgt_pts), dim=0) pose_to_frame1 = torch.cat((pose_to_frame, pose_to_frame), dim=0) srcs1 = torch.bmm(pose_to_frame1, srcs.permute(0, 2, 1)).permute(0, 2, 1) loss = torch.mean(torch.abs(srcs1[:, :, :2] - tgts[:, :, :2])) return loss def rand_dis(x, y): assert len(x.shape)==2 and len(y.shape)==2,'x and y must be 2 dim' N, N = x.size() ids=torch.arange(N).to(x.device) idx = ids.view(1, N).repeat(N, 1) mask = ~(idx == idx.transpose(0, 1)) idx1 = idx[mask].view(N, N - 1) random_indices = torch.randint(N - 1, size=(N,)).to(x.device) rand_idx = torch.gather(idx1, 1, random_indices.view(-1, 1)) rand_idx1 = torch.cat([ids.view(N, 1), rand_idx], dim=1) diag = ids.view(N, 1).repeat(1, 2) x1 = x[rand_idx1[:, 0], rand_idx1[:, 1]] x2 = x1*0 x3 = torch.cat([x1.view(N, 1), x2.view(N, 1)], dim=1) y1 = y[rand_idx1[:, 0], rand_idx1[:, 1]] y2 = y[diag[:, 0], diag[:, 1]] y3 = torch.cat([y1.view(N, 1), y2.view(N, 1)], dim=1) dis=torch.abs(x3-y3).mean()+F.relu(0.2-torch.abs(y1)).mean() return dis def gen_feature_loss(batch_dict): #BCN fea_pt_dual_gen = batch_dict['fea_pt_dual_gen'] fea_pl_dual_gen = batch_dict['fea_pl_dual_gen'] fea_kpt_original_gen = batch_dict['fea_kpt_original_gen'] # fea_kpt_gen_gen=batch_dict['fea_kpt_gen_gen'] fea_pt_dual = batch_dict['fea_pt_dual'] fea_pl_dual = batch_dict['fea_pl_dual'] fea_kpt_original = batch_dict['fea_kpt_original'] # fea_pt_dual = batch_dict['fea_pt_dual'].detach() # fea_pl_dual = batch_dict['fea_pl_dual'].detach() # fea_kpt_original = batch_dict['fea_kpt_original'].detach() b = fea_pl_dual.shape[0] loss0 = 0 loss1 = 0 loss2 = 0 loss3 = 0 relation = batch_dict['relation'] nums=0 for i in range(b): cnt = torch.sum((relation[i, :, -1, 0] > 0) & (relation[i, :, -1, 1] > 0)) nums+=cnt fea_pt_dual1 = fea_pt_dual[i, :, :cnt] # 匹配点云特征，CN fea_pt_dual_gen1 = fea_pt_dual_gen[i, :, :cnt] # 匹配点云特征，生成于图像 fea_pl_dual1 = fea_pl_dual[i, :, :cnt] # 匹配图像特征 fea_pl_dual_gen1 = fea_pl_dual_gen[i, :, :cnt] # 匹配图像特征，生成于点云 # loss0 = loss0 + torch.abs(fea_pt_dual1 - fea_pt_dual_gen1).mean() loss0 = loss0 + (1 - F.cosine_similarity(fea_pt_dual1,fea_pt_dual_gen1,dim=0)).mean() # loss0 = loss0 + F.mse_loss(fea_pt_dual1, fea_pt_dual_gen1) # loss0 = loss0 + ((fea_pt_dual1 - fea_pt_dual_gen1).norm(p=2, dim=0)).mean() # sims00=tools.batch_distance(fea_pt_dual1.unsqueeze(0).permute(0,2,1),fea_pt_dual1.unsqueeze(0).permute(0,2,1),'cosine') # sims01=tools.batch_distance(fea_pt_dual_gen1.unsqueeze(0).permute(0,2,1),fea_pt_dual_gen1.unsqueeze(0).permute(0,2,1),'cosine') # loss0 = loss0 + torch.abs(sims00-sims01).mean() # loss1 = loss1 + torch.abs(fea_pl_dual1 - fea_pl_dual_gen1).mean() loss1 = loss1 + (1 - F.cosine_similarity(fea_pl_dual1,fea_pl_dual_gen1,dim=0)).mean() # loss1 = loss1 + F.mse_loss(fea_pl_dual1, fea_pl_dual_gen1) # loss1 = loss1 + ((fea_pl_dual1 - fea_pl_dual_gen1).norm(p=2, dim=0)).mean() # sims10=tools.batch_distance(fea_pl_dual1.unsqueeze(0).permute(0,2,1),fea_pl_dual1.unsqueeze(0).permute(0,2,1),'cosine') # sims11=tools.batch_distance(fea_pl_dual_gen1.unsqueeze(0).permute(0,2,1),fea_pl_dual_gen1.unsqueeze(0).permute(0,2,1),'cosine') # loss1 = loss1 + torch.abs(sims10-sims11).mean() #全景特征生成模块损失计算 # loss2 = loss2 + torch.abs(fea_kpt_original[i] - fea_kpt_original_gen[i]).mean() loss2= loss2 + (1-F.cosine_similarity(fea_kpt_original[i], fea_kpt_original_gen[i],dim=0)).mean() # loss2 = loss2 + F.mse_loss(fea_kpt_original[i], fea_kpt_original_gen[i]) # loss2 = loss2 + ((fea_kpt_original[i] - fea_kpt_original_gen[i]).norm(p=2, dim=0)).mean() # sims20=tools.batch_distance(fea_kpt_original[i:i+1].permute(0,2,1),fea_kpt_original[i:i+1].permute(0,2,1),'cosine') # sims21=tools.batch_distance(fea_kpt_original_gen[i:i+1].permute(0,2,1),fea_kpt_original_gen[i:i+1].permute(0,2,1),'cosine') # loss2 = loss2 + torch.abs(sims20-sims21).mean() loss0 = loss0 / b loss1 = loss1 / b loss2 = loss2 / b return loss0, loss1, loss2, loss3 def sinkhorn_matches_loss(batch_dict,key): project_kpts = batch_dict[key] # calculated from corrspondence of kpts src_coords = batch_dict['key_points'] pose_to_frame = batch_dict['pose_to_frame'] src_coords = src_coords.clone().view(batch_dict['batch_size'], -1, 4) B, N_POINT, _ = src_coords.shape B = B // 2 src_coords = src_coords[:B, :, [0, 1, 2, 3]] src_coords[:, :, -1] = 1. gt_dst_coords = torch.bmm(pose_to_frame, src_coords.permute(0, 2, 1)) # True project kpts gt_dst_coords = gt_dst_coords.permute(0, 2, 1)[:, :, :3] loss = (gt_dst_coords - project_kpts).norm(dim=2).mean() return loss def score_loss(batch_dict): score = batch_dict['score_bev'] label_score = batch_dict['label_score'] label_score = torch.cat([label_score[:, :, :, 0], label_score[:, :, :, 1]], dim=0) mask1 = score > 1e-8 # mask2 = label_score > 1e-8 # mask = mask1 | mask2 score = score[mask1] label_score = label_score[mask1] loss = nn.functional.mse_loss(score, label_score) return loss def pose_loss(batch_dict,key): src_coords = batch_dict['key_points'] src_coords = src_coords.clone().view(batch_dict['batch_size'], -1, 4) delta_pose = batch_dict['pose_to_frame'] B, N_POINT, _ = src_coords.shape B = B // 2 src_coords = src_coords[:B] gt_dst_coords = torch.bmm(delta_pose, src_coords.permute(0, 2, 1)).float() gt_dst_coords = gt_dst_coords.permute(0, 2, 1)[:, :, :3] transformation = batch_dict[key] pred_dst_coords = torch.bmm(transformation, src_coords.permute(0, 2, 1)) pred_dst_coords = pred_dst_coords.permute(0, 2, 1)[:, :, :3] loss = torch.mean(torch.abs(pred_dst_coords - gt_dst_coords)) return loss def get_all_triplets(dist_mat, pos_mask, neg_mask, is_inverted=False, margin=0.5, different_embedding=False): if not different_embedding: pos_mask = torch.triu(pos_mask, 1) triplets = pos_mask.unsqueeze(2) * neg_mask.unsqueeze(1) return torch.where(triplets) def hardest_negative_selector(dist_mat, pos_mask, neg_mask, is_inverted, margin=0.5, different_embedding=False): if not different_embedding: pos_mask = torch.triu(pos_mask, 1) a, p = torch.where(pos_mask) if neg_mask.sum() == 0: return a, p, None if is_inverted: dist_neg = dist_mat * neg_mask n = torch.max(dist_neg, dim=1) else: dist_neg = dist_mat.clone() dist_neg[~neg_mask] = dist_neg.max() + 1. _, n = torch.min(dist_neg, dim=1) n = n[a] return a, p, n def random_negative_selector(dist_mat, pos_mask, neg_mask, is_inverted, margin=0.5, different_embedding=False): if not different_embedding: pos_mask = torch.triu(pos_mask, 1) a, p = torch.where(pos_mask) selected_negs = [] for i in range(a.shape[0]): possible_negs = torch.where(neg_mask[a[i]])[0] if len(possible_negs) == 0: return a, p, None dist_neg = dist_mat[a[i], possible_negs] if is_inverted: curr_loss = -dist_mat[a[i], p[i]] + dist_neg + margin else: curr_loss = dist_mat[a[i], p[i]] - dist_neg + margin if len(possible_negs[curr_loss > 0]) > 0: possible_negs = possible_negs[curr_loss > 0] random_neg = np.random.choice(possible_negs.cpu().numpy()) selected_negs.append(random_neg) n = torch.tensor(selected_negs, dtype=a.dtype, device=a.device) return a, p, n def semihard_negative_selector(dist_mat, pos_mask, neg_mask, is_inverted, margin=0.5, different_embedding=False): if not different_embedding: pos_mask = torch.triu(pos_mask, 1) a, p = torch.where(pos_mask) selected_negs = [] for i in range(a.shape[0]): possible_negs = torch.where(neg_mask[a[i]])[0] if len(possible_negs) == 0: return a, p, None dist_neg = dist_mat[a[i], possible_negs] if is_inverted: curr_loss = -dist_mat[a[i], p[i]] + dist_neg + margin else: curr_loss = dist_mat[a[i], p[i]] - dist_neg + margin semihard_idxs = (curr_loss > 0) & (curr_loss < margin) if len(possible_negs[semihard_idxs]) > 0: possible_negs = possible_negs[semihard_idxs] random_neg = np.random.choice(possible_negs.cpu().numpy()) selected_negs.append(random_neg) n = torch.tensor(selected_negs, dtype=a.dtype, device=a.device) return a, p, n class TripletLoss(nn.Module): def __init__(self, margin: float, triplet_selector, distance: distances.BaseDistance): super(TripletLoss, self).__init__() self.margin = margin self.triplet_selector = triplet_selector self.distance = distance def forward(self, embeddings, pos_mask, neg_mask, other_embeddings=None): if other_embeddings is None: other_embeddings = embeddings dist_mat = self.distance(embeddings, other_embeddings) triplets = self.triplet_selector( dist_mat, pos_mask, neg_mask, self.distance.is_inverted) distance_positive = dist_mat[triplets[0], triplets[1]] if triplets[-1] is None: if self.distance.is_inverted: return F.relu(1 - distance_positive).mean() else: return F.relu(distance_positive).mean() distance_negative = dist_mat[triplets[0], triplets[2]] curr_margin = self.distance.margin( distance_positive, distance_negative) loss = F.relu(curr_margin + self.margin) return loss.mean() def _pairwise_distance(x, squared=False, eps=1e-16): # Compute the 2D matrix of distances between all the embeddings. cor_mat = torch.matmul(x, x.t()) norm_mat = cor_mat.diag() distances = norm_mat.unsqueeze(1) - 2 * cor_mat + norm_mat.unsqueeze(0) distances = F.relu(distances) if not squared: mask = torch.eq(distances, 0.0).float() distances = distances + mask * eps distances = torch.sqrt(distances) distances = distances * (1.0 - mask) return distances class TotalLoss(nn.Module): def __init__(self, cfg): super(TotalLoss, self).__init__() if 'hardest' == cfg['negetative_selsector']: neg_selector = hardest_negative_selector elif 'semihard' == cfg['negetative_selsector']: neg_selector = semihard_negative_selector else: neg_selector = random_negative_selector self.trip_fun = TripletLoss(margin=cfg['trip_margin'], triplet_selector=neg_selector, distance=distances.LpDistance()) self.negetative_distcance = 50 def forward(self, batch_dict): l_pose=l_score=l_match=l_tra=l_rot=l_gb=l_gi=l_gpa=l_gpo=l_kpl = 0 if 'key_points' in batch_dict.keys(): l_score = score_loss(batch_dict) l_match1,l_pose1,l_match2,l_pose2,l_tra1,l_rot1,l_tra2,l_rot2=0,0,0,0,0,0,0,0 if 'transformation_original' in batch_dict.keys(): l_match1 = sinkhorn_matches_loss(batch_dict,'project_kpts_original') l_tra1, l_rot1 = tr_loss(batch_dict,'transformation_original') l_pose1 = pose_loss(batch_dict,'transformation_original') if 'transformation_fusion' in batch_dict.keys(): l_match2 = sinkhorn_matches_loss(batch_dict,'project_kpts_fusion') l_tra2, l_rot2 = tr_loss(batch_dict,'transformation_fusion') l_pose2 = pose_loss(batch_dict,'transformation_fusion') cnt=1 if min(l_rot1,l_rot2)>0: cnt=2 l_match=(l_match1+l_match2)/cnt l_pose=(l_pose1+l_pose2)/cnt l_tra=(l_tra1+l_tra2)/cnt l_rot=(l_rot1+l_rot2)/cnt if ('fea_pt_dual_gen' in batch_dict.keys()) or ('fea_pl_dual_gen' in batch_dict.keys()): l_gb, l_gi, l_gpa,l_kpl = gen_feature_loss(batch_dict) if 'key_points_gen' in batch_dict.keys(): l_gpo = gen_points_loss(batch_dict) if 'sequence' in batch_dict: neg_mask = batch_dict['sequence'].view(1, -1) != batch_dict['sequence'].view(-1, 1) else: neg_mask = torch.zeros((batch_dict['pose_query'].shape[0] * 2, batch_dict['pose_query'].shape[0] * 2), dtype=torch.bool) pair_dist = _pairwise_distance(batch_dict['pose_query'][:, 0:3, 3]) neg_mask = ((pair_dist > self.negetative_distcance) | neg_mask.to(pair_dist.device)) neg_mask = neg_mask.repeat(2, 2) batch_size = batch_dict['batch_size'] pos_mask = torch.zeros((batch_size, batch_size), device=neg_mask.device) for i in range(batch_size // 2): pos_mask[i, i + batch_size // 2] = 1 pos_mask[i + batch_size // 2, i] = 1 l_triplet = self.trip_fun(batch_dict['vlads'], pos_mask, neg_mask) l_total = l_score + l_pose + 0.05 * l_match + l_triplet + (l_gb + l_gi + l_gpa + l_kpl) loss = [l_total, l_pose, l_score, l_match, l_triplet, l_tra, l_rot, l_gb, l_gi, l_gpa, l_gpo,l_kpl] for i in range(len(loss)): if loss[i]==0: loss[i]=loss[0]*0 batch_dict['loss']=loss return batch_dict