[Python] GraphMAE self-supervised learning: node attribute (sampled points) reconstruction...

I am experimenting with a GraphMAE self-supervised architecture using PyTorch + DGL. In my task, each graph node represents a CAD entity, and one node attribute stores sampled points (coordinates + tangent + normal). The goal is to mask these attributes and let the model reconstruct them.

The problem

In theory, if I use a single graph as input, the model should learn to reconstruct sampled points correctly.

However, in my main training pipeline, the reconstructed sampled points are completely scattered (not matching the target shape), even though the loss decreases normally.

To debug, I wrote a simplified training script that uses the same data load/feature process/encoder/decoder architecture. → In this simplified version, the model can successfully reconstruct sampled points as expected.

I compared the full pipeline with the minimal script, printed lots of debugging information (shapes, ranges, mask logic, loss components), but I still cannot identify where the inconsistency comes from.

Question

What kinds of hidden differences in my code could cause GraphMAE to fail to reconstruct node attributes in one pipeline but succeed in another, even when:

encoder and decoder are the same,

loss decreases normally, and

debug info (shapes/ranges/masks) looks consistent?

How should I systematically debug to find where the pipelines diverge?

What I tried

Verified that target and prediction tensors have the same shape and index alignment (especially with masked vs unmasked nodes).

Checked value ranges of predicted vs target sampled points → both are within the same numerical range.

Debugged loss components (coords/tangent/normal): loss decreases in both cases.

Tested single-sample training (batch size = 1), but still scattered predictions in the full pipeline.

What I expected

I expected both the full pipeline and the simplified training script to produce similar reconstructions, since the architectures and losses are (in principle) the same.

here is my simplified training script:

def oracle_test_with_complete_pipeline(target_data, encoder_model, decoder_type='mlp', num_epochs=2000, lr=0.001, freeze_encoder=False, mask_ratio=0.75, mask_strategy='zero'):
...

for epoch in range(num_epochs):
optimizer.zero_grad()

try:
# === MAE-style Masking ===
# 每个epoch重新生成mask（可选：也可以固定mask）
mask_indices, keep_indices, _ = create_node_mask(num_nodes, mask_ratio, device)

# === Feature Processing ===
# 对图的节点和边特征进行预处理
try:
# 临时将masked特征设置到图中进行处理
original_node_feat = graph.ndata['feat'].clone()

# 应用mask到输入特征
x_masked = apply_node_mask(x, mask_indices, mask_strategy, mask_token)

# 将masked特征设置到图中（注意：这里假设x就是graph.ndata['feat']的一部分或全部）
# 如果x的维度与graph.ndata['feat']不同，需要相应调整
if x.shape == graph.ndata['feat'].shape:
graph.ndata['feat'] = x_masked
else:
# 如果维度不匹配，保持原有特征不变，只在后续处理中使用x_masked
print(f"[WARNING] Input x shape {x.shape} != graph node feat shape {graph.ndata['feat'].shape}")

# 应用特征处理器
processed_node_feat, processed_edge_feat = feature_processor(graph)

# 恢复原始图特征
graph.ndata['feat'] = original_node_feat

# 添加feature processing调试信息
if epoch == 0 or (epoch + 1) % 200 == 0:
print(f"[FEATURE PROCESS DEBUG] Epoch {epoch+1}:")
print(f" - Original node feat shape: {x.shape}")
print(f" - Processed node feat shape: {processed_node_feat.shape}")
print(f" - Processed edge feat shape: {processed_edge_feat.shape}")
print(f" - Processed node feat range: [{processed_node_feat.min().item():.6f}, {processed_node_feat.max().item():.6f}]")
print(f" - Processed edge feat range: [{processed_edge_feat.min().item():.6f}, {processed_edge_feat.max().item():.6f}]")

# 添加mask调试信息
if epoch == 0 or (epoch + 1) % 200 == 0:
print(f"[MASK DEBUG] Epoch {epoch+1}:")
print(f" - Original input range: [{x.min().item():.6f}, {x.max().item():.6f}]")
print(f" - Masked input range: [{x_masked.min().item():.6f}, {x_masked.max().item():.6f}]")
if len(mask_indices) > 0:
print(f" - Mask indices (first 10): {mask_indices[:10].tolist()}")
print(f" - Keep indices (first 10): {keep_indices[:10].tolist()}")

except Exception as fp_e:
print(f"[WARNING] Feature processing failed: {fp_e}")
print("Falling back to original features...")
# 如果特征处理失败，使用原始特征
x_masked = apply_node_mask(x, mask_indices, mask_strategy, mask_token)
processed_node_feat = x_masked
processed_edge_feat = edge_feat if edge_feat is not None else torch.zeros(graph.num_edges(), 64, device=device)

# 获取真正的编码器隐向量（使用处理后的特征）
# 对于PreModel，需要直接调用embed方法获取编码器输出，而不是使用mask_attr_prediction的解码结果
if hasattr(encoder_model, 'embed'):
# PreModel有embed方法，直接获取编码器隐向量
latent_code = encoder_model.embed(graph, processed_node_feat, edge_feat=processed_edge_feat).to(device)
if epoch == 0:
print(f"[DEBUG] Using encoder.embed() with processed features to get latent code, shape: {latent_code.shape}")
else:
# 其他编码器类型的处理
encoder_output = encoder_model(graph, processed_node_feat, edge_feat=processed_edge_feat)

# 获取latent code
if isinstance(encoder_output, dict):
if latent_key in encoder_output:
latent_code = encoder_output[latent_key].to(device)
else:
latent_code = encoder_output[list(encoder_output.keys())[0]].to(device)
elif isinstance(encoder_output, tuple):
# 处理元组输出 - 注意：这里不应该是PreModel的mask_attr_prediction结果
if len(encoder_output) == 3:
# MPNN编码器返回的(h_node, h_edge, h_global)
latent_code = encoder_output[0].to(device) # 使用节点特征
elif len(encoder_output) > 0:
latent_code = encoder_output[0].to(device) # 默认使用第一个元素
else:
raise ValueError("Empty tuple from encoder output")
else:
latent_code = encoder_output.to(device)

# 处理latent code维度
if len(latent_code.shape) == 3: # [1, num_nodes, latent_dim] 或 [batch_size, num_nodes, latent_dim]
latent_code = latent_code.squeeze(0) # [num_nodes, latent_dim]

# 确保latent_code在正确设备上
latent_code = latent_code.to(device)

# 创建target_fusion_features (使用处理后的节点特征作为目标)
# fusion decoder的目标应该是经过特征处理的节点特征，维度为fusion_dim=256
if target_fusion_features is None:
# 使用处理后的节点特征作为fusion目标，但需要调整到fusion_dim
# 这里我们使用一个简单的线性变换将processed_node_feat映射到fusion_dim
fusion_dim = 256
if processed_node_feat.shape[1] != fusion_dim:
# 创建一个临时的线性层来调整维度
temp_proj = nn.Linear(processed_node_feat.shape[1], fusion_dim).to(device)
with torch.no_grad():
target_fusion_features = temp_proj(processed_node_feat).detach()
else:
target_fusion_features = processed_node_feat.clone().detach()
print(f"Created target_fusion_features with shape: {target_fusion_features.shape}")

# Step 1: 通过fusion decoder得到融合特征
if decoder_type == 'mpnn':
# MPNN fusion decoder使用原始图结构
node_features = latent_code.to(device) # 确保在正确设备上

# 边特征：使用节点特征的平均值作为边特征
src, dst = graph.edges()
src, dst = src.to(device), dst.to(device) # 确保边索引在正确设备上
edge_features = (node_features[src] + node_features[dst]) / 2

print(f"[DEBUG] MPNN inputs:")
print(f" - node_features shape: {node_features.shape}")
print(f" - edge_features shape: {edge_features.shape}")

# 通过MPNN生成fusion features
fusion_features, _, _ = fusion_decoder(graph, node_features, edge_features)
fusion_features = fusion_features.to(device) # 确保输出在正确设备上
else:
# 使用MLP fusion decoder
fusion_features = fusion_decoder(latent_code.to(device))
fusion_features = fusion_features.to(device) # 确保输出在正确设备上

# Step 2: 使用fusion features进行多个解码器前向传播
pred_basic = basic_decoder(fusion_features) # [N, 7]
pred_bbox = bbox_decoder(fusion_features) # [N, 4]

# 采样点解码器
pred_tangent = None
pred_normal = None
if decoder_type == 'folding':
pred_points_output = points_decoder(fusion_features)
if isinstance(pred_points_output, dict):
pred_points = pred_points_output['sampled_points']
pred_tangent = pred_points_output.get('tangent_vectors', None)
pred_normal = pred_points_output.get('normal_vectors', None)
else:
pred_points = pred_points_output
else:
pred_points_flat = points_decoder(fusion_features) # [N, 128]
pred_points = pred_points_flat.view(num_nodes, 64, 2) # [N, 64, 2]

# === MAE-style Loss Calculation ===
# 只对被mask的节点计算重建损失
if len(mask_indices) > 0:
# 提取被mask节点的预测和目标
pred_basic_masked = pred_basic[mask_indices]
pred_bbox_masked = pred_bbox[mask_indices]
pred_points_masked = pred_points[mask_indices]

target_basic_masked = target_basic[mask_indices]
target_bbox_masked = target_bbox[mask_indices]
target_sampled_points_masked = target_sampled_points[mask_indices]

# 计算基础损失（只针对被mask的节点）
loss_basic = mse_loss(pred_basic_masked, target_basic_masked)
loss_bbox = mse_loss(pred_bbox_masked, target_bbox_masked)
loss_points = mse_loss(pred_points_masked, target_sampled_points_masked)

# 计算fusion decoder损失（只针对被mask的节点）
target_fusion_masked = target_fusion_features[mask_indices]
fusion_features_masked = fusion_features[mask_indices]
loss_fusion = mse_loss(fusion_features_masked, target_fusion_masked)

# 计算tangent和normal向量损失（如果可用）
loss_tangent = torch.tensor(0.0, device=device)
loss_normal = torch.tensor(0.0, device=device)

if decoder_type == 'folding' and pred_tangent is not None and target_tangent_vectors is not None:
pred_tangent_masked = pred_tangent[mask_indices]
target_tangent_masked = target_tangent_vectors[mask_indices]
loss_tangent = mse_loss(pred_tangent_masked, target_tangent_masked)

if decoder_type == 'folding' and pred_normal is not None and target_normal_vectors is not None:
pred_normal_masked = pred_normal[mask_indices]
target_normal_masked = target_normal_vectors[mask_indices]
loss_normal = mse_loss(pred_normal_masked, target_normal_masked)
else:
# 如果没有mask节点，使用全部节点计算损失（fallback）
loss_basic = mse_loss(pred_basic, target_basic)
loss_bbox = mse_loss(pred_bbox, target_bbox)
loss_points = mse_loss(pred_points, target_sampled_points)

# 计算fusion decoder损失（全部节点）
loss_fusion = mse_loss(fusion_features, target_fusion_features)

# 计算tangent和normal向量损失（如果可用）
loss_tangent = torch.tensor(0.0, device=device)
loss_normal = torch.tensor(0.0, device=device)

if decoder_type == 'folding' and pred_tangent is not None and target_tangent_vectors is not None:
loss_tangent = mse_loss(pred_tangent, target_tangent_vectors)

if decoder_type == 'folding' and pred_normal is not None and target_normal_vectors is not None:
loss_normal = mse_loss(pred_normal, target_normal_vectors)

if epoch == 0:
print(f"[LOSS DEBUG] No masked nodes, using full reconstruction loss")

# 总损失 (包含所有组件)
total_loss = loss_basic + loss_bbox + loss_points + loss_fusion + loss_tangent + loss_normal

# 反向传播
total_loss.backward()
optimizer.step()

# 记录损失
current_loss = total_loss.item()
loss_history.append({
'total': current_loss,
'basic': loss_basic.item(),
'bbox': loss_bbox.item(),
'points': loss_points.item(),
'fusion': loss_fusion.item(),
'tangent': loss_tangent.item(),
'normal': loss_normal.item()
})

if current_loss < best_loss:
best_loss = current_loss


and here is the mask_attr_prediction function from encoder.forward ( I'am pretty sure the problem happens here but the logical looks exactly like the code above)

def mask_attr_prediction(self, g, x, targets, epoch, drop_g1=None, drop_g2=None, edge_feat=None, x_fused_original=None, use_global_normalization=False):
# 初始化损失组件字典
loss_components = {}

if x_fused_original is None:
pass

use_x = x

graph_for_encoder, features_for_encoder, mask_indices_info = self.encoding_mask_noise(g, use_x, self.mask_rate)
mask_nodes, keep_nodes, token_nodes = mask_indices_info
# --- End Masking ---

# --- GNN Encoding ---
enc_out_feat = None
enc_out_edge_feat = None

if self._encoder_type == "mpnn":
if hasattr(self.encoder, 'forward') and 'h_edge_initial' in self.encoder.forward.__code__.co_varnames:
encoder_outputs = self.encoder(graph_for_encoder, features_for_encoder, edge_feat)
enc_out_feat = encoder_outputs[0]
enc_out_edge_feat = encoder_outputs[1]
else:
enc_out_feat = self.encoder(graph_for_encoder, features_for_encoder)
enc_out_edge_feat = edge_feat
else:
enc_out_feat = self.encoder(graph_for_encoder, features_for_encoder)
enc_out_edge_feat = edge_feat

encoder_rep = enc_out_feat.clone()

# === Debug: Print Encoder Latent Features ===
print(f"\n=== Encoder Latent Features Debug ===")
print(f"Encoder output shape: {encoder_rep.shape}")
print(f"Encoder output range: [{encoder_rep.min().item():.6f}, {encoder_rep.max().item():.6f}]")
print(f"Encoder output mean: {encoder_rep.mean().item():.6f}")
print(f"Encoder output std: {encoder_rep.std().item():.6f}")
if hasattr(graph_for_encoder, 'ndata') and 'is_virtual' in graph_for_encoder.ndata:
real_nodes_idx = torch.where(~graph_for_encoder.ndata['is_virtual'])[0]
if len(real_nodes_idx) > 0:
real_latent = encoder_rep[real_nodes_idx]
print(f"Real nodes latent shape: {real_latent.shape}")
print(f"Real nodes latent range: [{real_latent.min().item():.6f}, {real_latent.max().item():.6f}]")
print(f"======================================\n")



feat_rec_final = None
loss_intermediate = torch.tensor(0.0, device=encoder_rep.device)
x_fused_rec = None

# --- DUAL DECODER ARCHITECTURE ---
# 修复问题3: 避免重复调用fusion_feature_decoder，统一处理边特征投影
if enc_out_edge_feat is not None and enc_out_edge_feat.shape[-1] != self.fusion_dim:
enc_out_edge_feat = self.edge_feat_proj_for_decoder(enc_out_edge_feat)

# 1. Fusion Feature Reconstruction - 融合特征解码器
if self.bypass_fusion_decoder:
# 跳过融合解码器，直接使用编码器输出
x_fused_rec = encoder_rep
else:
# 根据解码器类型使用不同的调用方式
if self.decoder_type == 'simple_mlp':
# SimpleMLP解码器只需要1个参数: (x)
sffd_output_tuple = self.fusion_feature_decoder(encoder_rep)
elif self.decoder_type.lower() == "mpnn":
# MPNN解码器需要3个参数: (g, h_node, h_edge)
sffd_output_tuple = self.fusion_feature_decoder(graph_for_encoder, encoder_rep, enc_out_edge_feat)
else:
# GCN/GAT解码器只需要2个参数: (g, inputs)
# 为GAT解码器添加自环以避免0入度节点问题
if self.decoder_type.lower() == "gat":
graph_for_decoder = dgl.add_self_loop(graph_for_encoder)
else:
graph_for_decoder = graph_for_encoder
sffd_output_tuple = self.fusion_feature_decoder(graph_for_decoder, encoder_rep)

x_fused_rec = sffd_output_tuple[0] if isinstance(sffd_output_tuple, tuple) else sffd_output_tuple

if x_fused_rec is None:
raise RuntimeError("fusion_feature_decoder returned None fused features.")

# Calculate Intermediate Loss
if x_fused_original is not None and self.intermediate_loss_weight > 0:
if self.mask_rate == 0: # Reconstruct all original fused features
if self.use_virtual_node:
real_nodes_idx = torch.where(~graph_for_encoder.ndata['is_virtual'])[0]
target_for_intermediate_loss = x_fused_original
pred_for_intermediate_loss = x_fused_rec[real_nodes_idx]
else:
target_for_intermediate_loss = x_fused_original
pred_for_intermediate_loss = x_fused_rec
if pred_for_intermediate_loss.shape[0] != target_for_intermediate_loss.shape[0]:
print(f"[ERROR PreModel] Intermediate loss shape mismatch (all nodes). Pred count: {pred_for_intermediate_loss.shape[0]}, Target count: {target_for_intermediate_loss.shape[0]}")
else: # Reconstruct only masked original fused features (self.mask_rate > 0)
if len(mask_nodes) > 0:
target_for_intermediate_loss = x_fused_original[mask_nodes]
pred_for_intermediate_loss = x_fused_rec[mask_nodes]
else:
# 如果没有mask节点，跳过中间损失计算
loss_intermediate = torch.tensor(0.0, device=encoder_rep.device)
target_for_intermediate_loss = None
pred_for_intermediate_loss = None

if target_for_intermediate_loss is not None and pred_for_intermediate_loss is not None:
loss_intermediate = self.intermediate_criterion(pred_for_intermediate_loss, target_for_intermediate_loss)
# --- End Intermediate Fused Feature Reconstruction Loss ---

feat_rec_basic = None
feat_rec_bbox = None
sampled_points_rec = None

# Stage 2a: Basic Feature Reconstruction using MLPDecoder from x_fused_rec
feat_rec_basic_raw = self.basic_feature_decoder(x_fused_rec)
if feat_rec_basic_raw is None:
raise RuntimeError("basic_feature_decoder (MLPDecoder) returned None.")

# Stage 2b: Sampled Points Reconstruction using FoldingNetUVDecoder from x_fused_rec
sampled_points_rec = None
if self.reconstruct_sampled_points:
sampled_points_rec = self.sampled_points_decoder(x_fused_rec)

# Stage 2c: Bbox Reconstruction using Sequential MLP from x_fused_rec
feat_rec_bbox = self.bbox_decoder(x_fused_rec)

# feat_rec_basic_raw现在直接是7维基础特征，无需合并操作
feat_rec_basic = feat_rec_basic_raw # [N, 7] 基础特征: [type(1) + length(1) + convexity(3) + centroid(2)]

loss_feat, loss_bbox, loss_sampled_points = torch.tensor(0.0, device=encoder_rep.device), \
torch.tensor(0.0, device=encoder_rep.device), \
torch.tensor(0.0, device=encoder_rep.device)

if targets is not None:
for key, value in targets.items():
if isinstance(value, torch.Tensor):
targets[key] = value.to(encoder_rep.device)

target_features_raw = targets["feat"]

# if not hasattr(self, '_printed_target_stats'):
# print(f"DEBUG: target_features_raw stats: min={target_features_raw.min().item()}, max={target_features_raw.max().item()}, mean={target_features_raw.mean().item()}, std={target_features_raw.std().item()}")
# if target_features_raw.shape[0] > 0:
# for i in range(target_features_raw.shape[1]):
# print(f"DEBUG: target_features_raw dim {i} stats: min={target_features_raw[:, i].min().item()}, max={target_features_raw[:, i].max().item()}, mean={target_features_raw[:, i].mean().item()}, std={target_features_raw[:, i].std().item()}")
# self._printed_target_stats = True

if feat_rec_basic is not None:
# 从feat中提取基础特征数据（索引0-4和9-10，不包括bbox索引5-8）
# 基础特征包括：type(0) + length(1) + convexity(2-4) + centroid(9-10) = 7维
target_type = target_features_raw[:, 0:1] # 索引0：type
target_length = target_features_raw[:, 1:2] # 索引1：length
target_convexity = target_features_raw[:, 2:5] # 索引2-4：convexity
target_centroid = target_features_raw[:, 9:11] # 索引9-10：centroid
target_basic_raw = torch.cat([target_type, target_length, target_convexity, target_centroid], dim=1) # [N, 7]

# feat_rec_basic现在直接是7维基础特征，可以直接使用
target_basic_for_loss = target_basic_raw
pred_basic_for_loss = feat_rec_basic # 直接使用7维输出
if self.use_virtual_node and self.mask_rate == 0:
real_nodes_idx = torch.where(~g.ndata['is_virtual'])[0]
target_basic_for_loss = target_basic_raw[real_nodes_idx]
pred_basic_for_loss = feat_rec_basic[real_nodes_idx]
elif self.mask_rate > 0 and mask_nodes.numel() > 0:
target_basic_for_loss = target_basic_raw[mask_nodes]
pred_basic_for_loss = feat_rec_basic[mask_nodes]
elif self.mask_rate > 0 and mask_nodes.numel() == 0:
loss_feat = torch.tensor(0.0, device=encoder_rep.device)

if pred_basic_for_loss.shape[0] > 0:
loss_feat = self.feat_loss_fn(pred_basic_for_loss, target_basic_for_loss)

if self.reconstruct_bbox and feat_rec_bbox is not None:
# 从feat中提取bbox数据（索引5-8）
target_bbox_raw = target_features_raw[:, 5:9] # bbox在索引5-8位置
target_bbox_for_loss = target_bbox_raw
pred_bbox_for_loss = feat_rec_bbox
if self.use_virtual_node and self.mask_rate == 0:
real_nodes_idx = torch.where(~g.ndata['is_virtual'])[0]
target_bbox_for_loss = target_bbox_raw[real_nodes_idx]
pred_bbox_for_loss = feat_rec_bbox[real_nodes_idx]
elif self.mask_rate > 0 and mask_nodes.numel() > 0:
target_bbox_for_loss = target_bbox_raw[mask_nodes]
pred_bbox_for_loss = feat_rec_bbox[mask_nodes]
elif self.mask_rate > 0 and mask_nodes.numel() == 0:
loss_bbox = torch.tensor(0.0, device=encoder_rep.device)

if pred_bbox_for_loss.shape[0] > 0:
loss_bbox = self.bbox_loss_fn(pred_bbox_for_loss, target_bbox_for_loss)



# Sampled Points Loss Calculation
if self.reconstruct_sampled_points and sampled_points_rec is not None:
# 从feat中提取采样点数据（索引11+）
target_sampled_points = target_features_raw[:, 11:] # 采样点从索引11开始
target_points_for_loss = target_sampled_points

# 从sampled_points_rec字典中提取坐标、切向量、法向量
if isinstance(sampled_points_rec, dict):
pred_coords = sampled_points_rec['sampled_points'] # [N, num_points, 2]
pred_tangents = sampled_points_rec['tangent_vectors'] # [N, num_points, 2]
pred_normals = sampled_points_rec['normal_vectors'] # [N, num_points, 2]

# 组合成6维格式 [N, num_points, 6] -> [N, num_points*6]
pred_combined = torch.cat([pred_coords, pred_tangents, pred_normals], dim=-1) # [N, num_points, 6]
pred_sampled_points = pred_combined.view(pred_combined.shape[0], -1) # [N, num_points*6]
else:
pred_sampled_points = sampled_points_rec



pred_points_for_loss = pred_sampled_points
if self.use_virtual_node and self.mask_rate == 0:
real_nodes_idx = torch.where(~g.ndata['is_virtual'])[0]
target_points_for_loss = target_sampled_points[real_nodes_idx]
pred_points_for_loss = pred_sampled_points[real_nodes_idx]


elif self.mask_rate > 0 and mask_nodes.numel() > 0:
target_points_for_loss = target_sampled_points[mask_nodes]
pred_points_for_loss = pred_sampled_points[mask_nodes]
elif self.mask_rate > 0 and mask_nodes.numel() == 0:
loss_sampled_points = torch.tensor(0.0, device=encoder_rep.device)

if pred_points_for_loss.shape[0] > 0:
loss_sampled_points = self.sampled_points_loss_fn(pred_points_for_loss, target_points_for_loss)

# === 加权损失计算 ===
# 1. 融合特征重建损失权重
weighted_fusion_loss = self.intermediate_loss_weight * loss_intermediate

# 2. 基础属性损失权重
weighted_attributes_loss = self.final_rec_loss_weight * self.feat_weight * loss_feat

# 3. bbox损失权重
weighted_bbox_loss = self.final_rec_loss_weight * loss_bbox

# 4. Sampled Points损失权重
weighted_sampled_points_loss = self.final_rec_loss_weight * loss_sampled_points

loss_components = {
"loss_feat_basic": loss_feat.detach().item(),
"loss_intermediate_fused": loss_intermediate.detach().item(),
"loss_bbox_params": loss_bbox.detach().item(),
"loss_sampled_points": loss_sampled_points.detach().item(),
"weighted_basic_features": weighted_attributes_loss.detach().item(),
"weighted_intermediate_fused": weighted_fusion_loss.detach().item(),
"weighted_bbox_params": weighted_bbox_loss.detach().item(),
"weighted_sampled_points": weighted_sampled_points_loss.detach().item(),
}

# 总损失 = 融合特征损失 + 基础属性损失 + bbox损失 + 采样点损失
total_loss = weighted_fusion_loss + weighted_attributes_loss + weighted_bbox_loss + weighted_sampled_points_loss

if self.intermediate_loss_weight > 0:
loss_components["weighted_loss_intermediate_fused"] = weighted_fusion_loss.detach().item()

return total_loss, loss_components, feat_rec_basic, sampled_points_rec, enc_out_edge_feat, mask_nodes, feat_rec_bbox

Continue reading...