数据结构
data.x #(128, 16) 16 = triangles(9) + confidence(7)
"""
每个元素
[x1, y1, z1, x2, y2, z2, x3, y3, z3,
c_v1, c_v2, c_v3, c_e1, c_e2, c_e3, c_f]
"""
data.y #(128)encoded_x #（128, 576)
encoded_x_conv # (2, 96, 576)
decoded_x_conv # (2, 96, 4)
decoded_x # (128, 4)
Angle_loss 计算
\[L_{\text{angle}} = p_{\text{merge}}^{\alpha} \times \min_{t \in \{90, 180, 270, 360\}} | \text{Angle}_{\text{merge}} - t |
\]
\[Loss= w_1 L_{\text{classification}} + w_2 L_{\text{angle}}
\]
1、朴素实现
data.x #(num_triangles, 16) , 16 = triangles(9) + confidence(7)
# 每个元素
[x1, y1, z1, x2, y2, z2, x3, y3, z3,
c_v1, c_v2, c_v3, c_e1, c_e2, c_e3, c_f]data.y #(128)
# 每个 triangle 的信息（添加了角）
data.x_with_angle
{face_idx: i,vertices: [(x1, y1, z1), (x2, y2, z2), (x3, y3, z3)],angles: [α1, α2, α3]
}# 获取所有 room triangles#取它们的顶点(去重)(x1, y1, z1): {vertices: (x1, y1, z1),faces_idx: [1, 2, 3, ...], # 包含该顶点的room面idxangle: [α1, α2, α3, ...], # 该顶点在该面的angledecoded_x: [[p_11, p_12, p_13, p_14], [p_21, p_22, p_23, p_24], [p_31, p_32, p_33, p_34], ...]
}# 计算angle_merge, p_merge# loss 计算朴素实现
def angle_loss_torch(self, data, decoded_x):"""doc: 计算角度正则化损失return: total_angle_loss: 角度正则化损失"""# 1. 选出所有最大分类概率为room的三角形(及其索引)softmax_decoded_x = torch.nn.functional.softmax(decoded_x, dim=1) # 使用softmax得到不同分类的概率max_probs, max_indices = torch.max(softmax_decoded_x, dim=1)room_indices = torch.where(max_indices == 2)[0]# 获取每个面的三个角 data.x_with_anglevertices = data.x[:, :9].reshape(-1, 3, 3) # 提取顶点坐标(num_triangle, 3, 3)angles = angle_func(vertices) # 计算内角 (num_triangle, 3)data.x_with_angle = [{'triangle_idx': i,'vertices': [tuple(v) for v in vertices[i]],'angles': angles[i].tolist()}for i in range(vertices.shape[0])]# 2. 取它们的顶点(去重)vertex_info = {}for i in room_indices:triangle = data.x_with_angle[i]face_idx = triangle['triangle_idx']vertices = triangle['vertices']angles = triangle['angles']for vertex, angle in zip(vertices, angles):vertex_key = tuple(vertex)if vertex_key not in vertex_info:vertex_info[vertex_key] = {'vertices': vertex,'faces_idx': [],'angles': [],'decoded_x': []}vertex_info[vertex_key]['faces_idx'].append(face_idx)vertex_info[vertex_key]['angles'].append(angle)vertex_info[vertex_key]['decoded_x'].append(softmax_decoded_x[face_idx])# 3. 计算每个顶点的角度总和和融合概率特征for vertex_key in vertex_info:vertex_info[vertex_key]['angle_merge'] = sum(vertex_info[vertex_key]['angles'])vertex_info[vertex_key]['p_merge'] = torch.mean(torch.stack(vertex_info[vertex_key]['decoded_x']), dim=0)# 4. angle_loss = p^alpha * min|θ - t|t_values = torch.tensor([np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], dtype=torch.float32)alpha = self.config.angle_loss_alphatotal_angle_loss = 0for vertex_key, info in vertex_info.items():angle_sum = info['angle_merge']p = info['p_merge'][2]abs_diffs = torch.abs(angle_sum - t_values)min_abs_diff = torch.min(abs_diffs)angle_loss = p**alpha * min_abs_difftotal_angle_loss += angle_lossreturn total_angle_loss
for循环太多，考虑能不能优化
2、优化循环
算法优化
# room 顶点
room_vertices = vertices[room_indices].reshape(-1, 3) # room angle
room_angles = angles[room_indices].reshape(-1) # (num_room_triangles * 3)# 从 softmax_decoded_x 中提取出属于“room”类别的三角形的概率分布，并将每个三角形的概率分布重复 3 次，以对应每个三角形的 3 个顶点
room_decoded_x = softmax_decoded_x[room_indices].repeat_interleave(3, dim=0) # (num_room_triangles * 3, num_cls)# 获取去重顶点
unique_vertices
# room_vertices 中相应元素在 unique_vertices 中的索引
inverse_indeces# angle_merge、p_merge
angle_merge.scatter_add_(0, inverse_indices, room_angles)
优化循环
def angle_loss(self, data, decoded_x):"""doc: 计算角度正则化损失return: total_angle_loss: 角度正则化损失"""# 1. 选出所有最大分类概率为room的三角形(及其索引)softmax_decoded_x = torch.nn.functional.softmax(decoded_x, dim=1) # 使用softmax得到不同分类的概率max_probs, max_indices = torch.max(softmax_decoded_x, dim=1)room_indices = torch.where(max_indices == 2)[0]# 获取每个面的三个角 data.x_with_anglevertices = data.x[:, :9].reshape(-1, 3, 3) # 提取顶点坐标(num_triangle, 3, 3)angles = angle_func(vertices) # 计算内角 (num_triangle, 3)# 2. 取它们的顶点(去重)room_vertices = vertices[room_indices].reshape(-1, 3) # (num_room_triangles * 3, 3)room_angles = angles[room_indices].reshape(-1) # (num_room_triangles * 3)room_decoded_x = softmax_decoded_x[room_indices].repeat_interleave(3, dim=0) # (num_room_triangles * 3, num_cls)# 3. 计算每个顶点的角度总和和融合概率特征unique_vertices, inverse_indices = torch.unique(room_vertices, dim=0, return_inverse=True)angle_merge = torch.zeros(unique_vertices.shape[0], dtype=torch.float32, device=room_vertices.device)p_merge = torch.zeros(unique_vertices.shape[0], softmax_decoded_x.shape[1], dtype=torch.float32, device=room_vertices.device)angle_merge.scatter_add_(0, inverse_indices, room_angles)p_merge.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, softmax_decoded_x.shape[1]), room_decoded_x)p_merge = p_merge / torch.bincount(inverse_indices).unsqueeze(1).float()# 4. angle_loss = p^alpha * min|θ - t|t_values = torch.tensor([np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], dtype=torch.float32, device=angle_merge.device)alpha = self.config.angle_loss_alphatotal_angle_loss = 0for angle_sum, p in zip(angle_merge, p_merge):abs_diffs = torch.abs(angle_sum - t_values)min_abs_diff = torch.min(abs_diffs)angle_loss = p[2]**alpha * min_abs_difftotal_angle_loss += angle_lossreturn total_angle_loss
