数据结构
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面idx
angle: [α1, α2, α3, ...], # 该顶点在该面的angle
decoded_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_angle
vertices = 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_alpha
total_angle_loss = 0
for 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_diff
total_angle_loss += angle_loss
return 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_angle
vertices = 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_alpha
total_angle_loss = 0
for 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_diff
total_angle_loss += angle_loss
return total_angle_loss