import os import sys import time project_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) sys.path.append(project_dir) os.chdir(project_dir) import torch from model.network import gradient class LossManager: def __init__(self, ablation, **condition_kwargs): self.weights = { "manifold": 1, "feature_manifold": 1, # 原文里面和manifold的权重是一样的 "normals": 1, "eikonal": 1, "offsurface": 1, "consistency": 1, "correction": 1, } self.condition_kwargs = condition_kwargs self.ablation = ablation # 消融实验用 def _get_condition_kwargs(self, key): """ 获取条件参数, 期望 ab: 损失类型 【overall, patch, off, cons, cc, cor,】 siren: 是否使用SIREN epoch: 当前epoch baseline: 是否为baseline """ if key in self.condition_kwargs: return self.condition_kwargs[key] else: raise ValueError(f"Key {key} not found in condition_kwargs") def pre_process(self, mnfld_pnts, mnfld_pred_all, nonmnfld_pnts, nonmnfld_pred_all, n_batchsize, n_branch, n_patch_batch, n_patch_last): """ 预处理 """ mnfld_pred_h = mnfld_pred_all[:,0] # 提取流形预测结果 nonmnfld_pred_h = nonmnfld_pred_all[:,0] # 提取非流形预测结果 mnfld_grad = gradient(mnfld_pnts, mnfld_pred_h) # 计算流形点的梯度 all_fi = torch.zeros([n_batchsize, 1], device = 'cuda') # 初始化所有流形预测值 for i in range(n_branch - 1): all_fi[i * n_patch_batch : (i + 1) * n_patch_batch, 0] = mnfld_pred_all[i * n_patch_batch : (i + 1) * n_patch_batch, i + 1] # 填充流形预测值 # last patch all_fi[(n_branch - 1) * n_patch_batch:, 0] = mnfld_pred_all[(n_branch - 1) * n_patch_batch:, n_branch] # 填充最后一个分支的流形预测值 return mnfld_pred_h, nonmnfld_pred_h, mnfld_grad, all_fi def position_loss(self, outputs): """ 计算流型损失的逻辑 :param outputs: 模型的输出 :return: 计算得到的流型损失值 """ # 计算流型损失（这里使用均方误差作为示例） manifold_loss = (outputs.abs()).mean() # 计算流型损失 return manifold_loss def normals_loss(self, normals: torch.Tensor, mnfld_pnts: torch.Tensor, all_fi: torch.Tensor, patch_sup: bool = True) -> torch.Tensor: """ 计算法线损失 :param normals: 法线 :param mnfld_pnts: 流型点 :param all_fi: 所有流型预测值 :param patch_sup: 是否支持补丁 :return: 计算得到的法线损失 """ # NOTE 源代码 这里还有复杂逻辑 # 计算分支梯度 branch_grad = gradient(mnfld_pnts, all_fi[:, 0]) # 计算分支梯度 # 计算法线损失 normals_loss = (((branch_grad - normals).abs()).norm(2, dim=1)).mean() # 计算法线损失 return normals_loss # 返回加权后的法线损失 def eikonal_loss(self, nonmnfld_pnts, nonmnfld_pred): """ 计算Eikonal损失 """ grad_loss_h = torch.zeros(1).cuda() # 初始化 Eikonal 损失 single_nonmnfld_grad = gradient(nonmnfld_pnts, nonmnfld_pred) # 计算非流形点的梯度 eikonal_loss = ((single_nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean() # 计算 Eikonal 损失 return eikonal_loss def offsurface_loss(self, nonmnfld_pnts, nonmnfld_pred): """ Eo 惩罚远离表面但是预测值接近0的点 """ offsurface_loss = torch.zeros(1).cuda() if not self.ablation == 'off': offsurface_loss = torch.exp(-100.0 * torch.abs(nonmnfld_pred)).mean() # 计算离表面损失 return offsurface_loss def consistency_loss(self, mnfld_pnts, mnfld_pred, all_fi): """ 惩罚流形点预测值和非流形点预测值不一致的点 """ mnfld_consistency_loss = torch.zeros(1).cuda() if not (self.ablation == 'cons' or self.ablation == 'cc'): mnfld_consistency_loss = (mnfld_pred - all_fi[:,0]).abs().mean() # 计算流形一致性损失 return mnfld_consistency_loss def correction_loss(self, mnfld_pnts, mnfld_pred, all_fi, th_closeness = 1e-5, a_correction = 100): """ 修正损失 """ correction_loss = torch.zeros(1).cuda() # 初始化修正损失 if not (self.ablation == 'cor' or self.ablation == 'cc'): mismatch_id = torch.abs(mnfld_pred - all_fi[:,0]) > th_closeness # 计算不匹配的 ID if mismatch_id.sum() != 0: # 如果存在不匹配 correction_loss = (a_correction * torch.abs(mnfld_pred - all_fi[:,0])[mismatch_id]).mean() # 计算修正损失 return correction_loss def compute_loss(self, mnfld_pnts, normals, mnfld_pred_all, nonmnfld_pnts, nonmnfld_pred_all, n_batchsize, n_branch, n_patch_batch, n_patch_last): """ 计算流型损失的逻辑 :param outputs: 模型的输出 :return: 计算得到的流型损失值 """ mnfld_pred, nonmnfld_pred, mnfld_grad, all_fi = self.pre_process(mnfld_pnts, mnfld_pred_all, nonmnfld_pnts, nonmnfld_pred_all, n_batchsize, n_branch, n_patch_batch, n_patch_last) manifold_loss = torch.zeros(1).cuda() # 计算流型损失（这里使用均方误差作为示例） if not self.ablation == 'overall': manifold_loss = (mnfld_pred.abs()).mean() # 计算流型损失 ''' if args.feature_sample: # 如果启用了特征采样 feature_indices = torch.randperm(args.all_feature_sample)[:args.num_feature_sample].cuda() # 随机选择特征点 feature_pnts = self.feature_data[feature_indices] # 获取特征点数据 feature_mask_pair = self.feature_data_mask_pair[feature_indices] # 获取特征掩码对 feature_pred_all = self.network(feature_pnts) # 进行前向传播，计算特征点的预测值 feature_pred = feature_pred_all[:,0] # 提取特征预测结果 feature_mnfld_loss = feature_pred.abs().mean() # 计算特征流形损失 loss = loss + weight_mnfld_h * feature_mnfld_loss # 将特征流形损失加权到总损失中 # patch loss: feature_id_left = [list(range(args.num_feature_sample)), feature_mask_pair[:,0].tolist()] # 获取左侧特征 ID feature_id_right = [list(range(args.num_feature_sample)), feature_mask_pair[:,1].tolist()] # 获取右侧特征 ID feature_fis_left = feature_pred_all[feature_id_left] # 获取左侧特征预测值 feature_fis_right = feature_pred_all[feature_id_right] # 获取右侧特征预测值 feature_loss_patch = feature_fis_left.abs().mean() + feature_fis_right.abs().mean() # 计算补丁损失 loss += feature_loss_patch # 将补丁损失加权到总损失中 # consistency loss: feature_loss_cons = (feature_fis_left - feature_pred).abs().mean() + (feature_fis_right - feature_pred).abs().mean() # 计算一致性损失 ''' manifold_loss_patch = torch.zeros(1).cuda() if self.ablation == 'patch': manifold_loss_patch = all_fi[:,0].abs().mean() # 计算法线损失 normals_loss = self.normals_loss(normals, mnfld_pnts, all_fi, patch_sup=True) # 计算Eikonal损失 eikonal_loss = self.eikonal_loss(nonmnfld_pnts, nonmnfld_pred_all) # 计算离表面损失 offsurface_loss = self.offsurface_loss(nonmnfld_pnts, nonmnfld_pred_all) # 计算一致性损失 consistency_loss = self.consistency_loss(mnfld_pnts, mnfld_pred, all_fi) # 计算修正损失 correction_loss = self.correction_loss(mnfld_pnts, mnfld_pred, all_fi) loss_details = { "manifold": self.weights["manifold"] * manifold_loss, "manifold_patch": manifold_loss_patch, "normals": self.weights["normals"] * normals_loss, "eikonal": self.weights["eikonal"] * eikonal_loss, "offsurface": self.weights["offsurface"] * offsurface_loss, "consistency": self.weights["consistency"] * consistency_loss, "correction": self.weights["correction"] * correction_loss, } # 计算总损失 total_loss = sum(loss_details.values()) return total_loss, loss_details