commit dbcb20280aca3f581cb951bccd18deb4c75e35ce Author: zhangzexuan <2944879106@qq.com> Date: Mon May 25 13:31:26 2026 +0800 上传文件至 / diff --git a/build_graph.py b/build_graph.py new file mode 100644 index 0000000..e4ab4c7 --- /dev/null +++ b/build_graph.py @@ -0,0 +1,217 @@ +import torch +from torch_geometric.data import Data +import numpy as np +import scipy.sparse as sp +import os +import glob +import re + +# ================= 配置区域 ================= +# 从全局配置文件导入数据集配置 +from config import DATASET_TYPE, DATA_ROOT_PATH, SCA_PREFIX, DATASET_DIRS_PATTERN + +# 根据全局配置设置路径 +root_data_path = os.path.join(os.path.dirname(__file__), DATA_ROOT_PATH) +sca_prefix = SCA_PREFIX +dataset_dirs_pattern = DATASET_DIRS_PATTERN +# =========================================== + +def is_main_process(): + """检查是否是主进程（rank 0）""" + rank = int(os.environ.get("RANK", "0")) + return rank == 0 + +# 全局变量 +data_mapping = {} +n_total = 0 + +# 移除模块级别的打印，避免DDP重复打印 +# 配置信息将在实际使用时打印 +# =========================================== + +def scan_all_data(root_path): + """ + 扫描A-4目录下的所有四组数据（A-TainDataset, A-TainDataset2, A-TainDataset3, A-TainDataset4） + 依然以 edge 文件为锚点 + """ + global data_mapping, n_total + + # 如果已经扫描过，直接返回，避免重复扫描 + if data_mapping and n_total > 0: + return data_mapping, n_total + + data_mapping = {} + k_idx = 1 + + # 获取所有数据子目录 + dataset_dirs = [d for d in os.listdir(root_path) + if os.path.isdir(os.path.join(root_path, d)) and d.startswith(dataset_dirs_pattern.replace('*', ''))] + + if is_main_process(): + print(f"发现 {len(dataset_dirs)} 个数据集目录: {dataset_dirs}") + + total_folders = 0 + for dataset_dir in sorted(dataset_dirs): + dataset_path = os.path.join(root_path, dataset_dir) + + # 扫描每个数据集目录中的sca*_data文件夹 + all_items = os.listdir(dataset_path) + subfolders = [f for f in all_items + if os.path.isdir(os.path.join(dataset_path, f)) + and f.startswith('sca') and f.endswith('_data')] + + def extract_folder_num(folder_name): + match = re.match(rf'{sca_prefix}(\d+)_data', folder_name) + return int(match.group(1)) if match else 9999 + + subfolders.sort(key=extract_folder_num) + + # 移除此打印，避免DDP重复打印 + + for folder_name in subfolders: + folder_path = os.path.join(dataset_path, folder_name) + folder_num = extract_folder_num(folder_name) + + # 查找 edge 文件 + search_pattern = os.path.join(folder_path, f"edge_{sca_prefix}{folder_num}_*.txt") + edge_files = glob.glob(search_pattern) + + file_ids = [] + pattern = re.compile(rf'edge_{sca_prefix}{folder_num}_(\d+)\.txt') + + for vf in edge_files: + match = pattern.match(os.path.basename(vf)) + if match: + file_ids.append(int(match.group(1))) + + file_ids.sort() + + for data_id in file_ids: + data_mapping[k_idx] = (folder_path, folder_num, str(data_id)) + k_idx += 1 + + total_folders += 1 + + n_total = k_idx - 1 + if is_main_process(): + print(f"初始化完成：总共扫描 {total_folders} 个文件夹，找到 {n_total} 组数据。") + return data_mapping, n_total + +def load_file_data(folder_path, prefix, folder_num, data_id): + """通用数据读取函数""" + filename = f"{prefix}_{sca_prefix}{folder_num}_{data_id}.txt" + filepath = os.path.join(folder_path, filename) + + if not os.path.exists(filepath): + raise FileNotFoundError(f"文件缺失: {filepath}") + + data = np.loadtxt(filepath, dtype=np.float32) + if data.ndim == 1: + data = data.reshape(1, -1) + return data + +def build_graph_data(k): + """ + 构建图数据，包含残差计算步骤: r = b - A * Ebz + """ + if not data_mapping: + scan_all_data(root_data_path) + + if k not in data_mapping: + raise ValueError(f"索引 k={k} 超出范围") + + folder_path, folder_num, data_id = data_mapping[k] + + # ------------------------------------------------------- + # 1. 基础数据读取 + # ------------------------------------------------------- + # 读取 edge (拓扑) + edge_data = load_file_data(folder_path, "edge", folder_num, data_id) + edge_index_np = edge_data.astype(np.int64) - 1 # MATLAB -> Python 索引 + edge_index = torch.tensor(edge_index_np.T, dtype=torch.long) + num_nodes = int(edge_index.max()) + 1 + + # 读取 eps (材料), b (右端项), Ebz (当前解/初始解) + eps_data = load_file_data(folder_path, "eps", folder_num, data_id) # [N, 2] + b_data = load_file_data(folder_path, "b", folder_num, data_id) # [N, 2] + Ebz_data = load_file_data(folder_path, "Ebz", folder_num, data_id) # [N, 2] + + # ------------------------------------------------------- + # 2. 计算残差 r (Preprocessing) + # 公式: r = b - A * Ebz + # ------------------------------------------------------- + + # 2.1 读取稀疏矩阵 A 的信息 (Aij 和 Av) + try: + Aij_data = load_file_data(folder_path, "Aij", folder_num, data_id) # [NNZ, 2] 坐标 + Av_data = load_file_data(folder_path, "Av", folder_num, data_id) # [NNZ, 2] 值 + except FileNotFoundError: + raise FileNotFoundError(f"计算残差需要 Aij 和 Av 文件，但在数据组 {k} 中未找到。") + + # 2.2 构建 Scipy 稀疏矩阵 (COO 格式) + # Aij 是 MATLAB 索引 (1-based)，需要减 1 + rows = Aij_data[:, 0].astype(int) - 1 + cols = Aij_data[:, 1].astype(int) - 1 + + # 构造复数数值: Real + 1j * Imag + values = Av_data[:, 0] + 1j * Av_data[:, 1] + + # 创建稀疏矩阵 A (N x N) + A_mat = sp.coo_matrix((values, (rows, cols)), shape=(num_nodes, num_nodes)) + + # 2.3 准备向量 b 和 Ebz (复数形式) + b_vec = b_data[:, 0] + 1j * b_data[:, 1] + Ebz_vec = Ebz_data[:, 0] + 1j * Ebz_data[:, 1] + + # 2.4 执行矩阵乘法和减法: r = b - A * Ebz + # A_mat.dot() 是高效的稀疏矩阵乘法 + Ax = A_mat.dot(Ebz_vec) + r_vec = b_vec - Ax # 得到复数残差向量 [N,] + + # 2.5 将残差拆分为实部和虚部 [N, 2] + r_real = r_vec.real + r_imag = r_vec.imag + + # ------------------------------------------------------- + # 3. 拼接节点特征 (Input Features) + # 目标输入: [eps, r, Ebz] (共 6 通道) + # ------------------------------------------------------- + # eps: [N, 2] + # r: [N, 2] (由上一步计算得到) + # Ebz: [N, 2] + + x_tensor = torch.cat([ + torch.from_numpy(eps_data), # eps_real, eps_imag + torch.tensor(r_real, dtype=torch.float32).unsqueeze(1), # r_real + torch.tensor(r_imag, dtype=torch.float32).unsqueeze(1), # r_imag + torch.from_numpy(Ebz_data), # Ebz_real, Ebz_imag (当前解) + torch.from_numpy(Ebz_data), # bg_real, bg_imag (背景场，不随网络更新) + ], dim=1) + + # ------------------------------------------------------- + # 4. 读取标签 (Ground Truth) + # ------------------------------------------------------- + Esz_data = load_file_data(folder_path, "Esz", folder_num, data_id) + y_tensor = torch.from_numpy(Esz_data) + + # ------------------------------------------------------- + # 5. 返回 Data 对象 + # ------------------------------------------------------- + data = Data(x=x_tensor, edge_index=edge_index, y=y_tensor) + data.k_idx = torch.tensor([k]) + + return data + +# 测试运行 +if __name__ == "__main__": + scan_all_data(root_data_path) + if n_total > 0: + print("\n--- 读取并计算残差中 ---") + try: + data = build_graph_data(1) + print("数据构建成功!") + print(f"节点特征 x shape: {data.x.shape}") + print("特征顺序: [eps_re, eps_im, r_re, r_im, Ebz_re, Ebz_im, bg_re, bg_im]") + print(f"残差(r)部分均值: {data.x[:, 2:4].abs().mean().item():.4e}") + except Exception as e: + print(f"出错: {e}") \ No newline at end of file diff --git a/config.py b/config.py new file mode 100644 index 0000000..e83d6ad --- /dev/null +++ b/config.py @@ -0,0 +1,110 @@ +""" +PhiSAGE 项目全局配置文件 +集中管理所有可配置的参数，避免循环导入问题 +""" + +# ========================================== +# 全局迭代次数配置（统一管理所有相关文件的迭代次数） +# ========================================== +N_ITER = 5 # 全局迭代次数，修改此值将影响所有相关文件 + +# ========================================== +# 全局模型保存路径配置 +# ========================================== +SAVE_DIR = "/public/home/zzx/gnn/PhiSAGE/PhiSAGE-test/saved_models/A-0.6" # 模型保存目录 +OUTPUT_DIR = "/public/home/zzx/gnn/PhiSAGE/PhiSAGE-test/training_outputs/A-0.6" # 训练曲线输出目录 +# ========================================== +# 数据集配置 +# ========================================== +# 数据集选择：设置为 'A' 或 'B' +# - 'A': 读取 data/A-4 目录（包含A-TainDataset等4个子目录），使用 scaA 命名格式 +# - 'B': 读取 data/B-4 目录（包含B-TainDataset等4个子目录），使用 scaB 命名格式 +# - 'C': 读取 data/C-4 目录（包含C-TainDataset等4个子目录），使用 scaC 命名格式 +# 使用方法： +# 1. 将数据放在对应的目录结构中 +# 2. 修改 DATASET_TYPE 的值 +# 3. 程序会自动扫描所有子目录中的数据 +# +DATASET_TYPE = 'A' # 可选: 'A' 或 'B' 或 'C' +DATA_ROOT_PATH = "data/A-1" + +# 根据选择动态设置相关参数（由 build_graph.py 使用） +if DATASET_TYPE == 'A': + SCA_PREFIX = 'scaA' + DATASET_DIRS_PATTERN = 'A-TainDataset*' +elif DATASET_TYPE == 'B': + SCA_PREFIX = 'scaB' + DATASET_DIRS_PATTERN = 'B-TainDataset*' +elif DATASET_TYPE == 'C': + SCA_PREFIX = 'scaC' + DATASET_DIRS_PATTERN = 'C-TainDataset*' +else: + raise ValueError("DATASET_TYPE 必须设置为 'A' 或 'B' 或 'C'") +# ========================================== +# ========================================== +# 训练启动器配置 (run.py 使用) +# ========================================== +# 默认 GPU 配置 +DEFAULT_TARGET_GPUS = "4,5" # 默认使用的 GPU 列表 + +# 训练模式配置 +DEFAULT_TRAIN_MODE = "ddp" # 默认训练模式: "single", "multi", "ddp" + +# 单卡训练配置 +DEFAULT_SINGLE_GPU_ID = 0 # 默认单卡 GPU ID + +# ========================================== +# DDP分布式训练配置 +# ========================================== +# DDP通信端口配置 +MASTER_PORT = "20870" # DDP主进程通信端口 + +# ========================================== +# 网络维度配置 +# ========================================== +# 网络架构配置：支持多种维度设置方式 +# +# 方式1：使用base_dim自动计算（推荐） +# NETWORK_BASE_DIM = 56 # 基础维度，会自动生成 [56, 112, 192] +# +# 方式2：自定义维度配置 +# NETWORK_CUSTOM_DIMS = [56, 112, 192] # 手动指定每层维度 +# +# 方式3：使用预设配置 +# NETWORK_CONFIG = "default" # 可选: "small", "medium", "large", "xlarge" + +# 当前使用的配置方式 +NETWORK_USE_CUSTOM_DIMS = True # True=使用自定义维度，False=使用base_dim自动计算 +NETWORK_BASE_DIM = 64 # 基础维度（当不使用自定义维度时） +NETWORK_CUSTOM_DIMS = [64, 128, 256] # 自定义维度配置 + +# 池化配置 +NETWORK_POOL_RATIOS = [0.8, 0.6] # 池化比例 + +# ========================================== +# Loss函数配置 +# ========================================== +# LOSS_TYPE: 选择使用的损失函数类型 +# - "mse": 传统的MSE损失 (||pred - true||^2) +# - "phi": Phi损失 (||A*x - b||^2) - 直接计算物理残差 +# - "asinh": Asinh损失 sqrt(asinh(norm(x-x_ref)^2/N)) - 平滑的损失函数 +# - "hybrid": 混合损失 - 前100epoch纯MSE，后续MSE + λ*phi (λ从0.001开始每50epoch×10倍，至0.1) +LOSS_TYPE = "hybrid" # 可选: "mse", "phi", "asinh" 或 "hybrid" + +# ========================================== +# 预训练模型加载配置（迁移学习） +# ========================================== +# LOAD_PRETRAINED_MODEL: 是否加载预训练模型权重（仅模型权重，不加载优化器等状态） +# - True: 从 PRETRAINED_MODEL_DIR 目录加载模型权重，优化器等重新初始化 +# - False: 不加载预训练模型，从头开始训练 +LOAD_PRETRAINED_MODEL = False # 是否加载预训练模型权重 + +# PRETRAINED_MODEL_DIR: 预训练模型目录 +# 如果 LOAD_PRETRAINED_MODEL=True，将从该目录加载模型权重文件（real_iter_*.pth 和 imag_iter_*.pth） +# None表示使用SAVE_DIR，也可以指定其他目录 +PRETRAINED_MODEL_DIR = None # None表示使用SAVE_DIR，也可以指定其他目录 + +# ========================================== +# 其他全局配置参数 +# ========================================== +# 可以在这里添加其他需要全局共享的配置参数 \ No newline at end of file diff --git a/model.py b/model.py new file mode 100644 index 0000000..fc647a9 --- /dev/null +++ b/model.py @@ -0,0 +1,467 @@ +import torch +import torch.nn as nn +from torch_geometric.nn import MessagePassing, TopKPooling + +# 从全局配置文件导入迭代次数配置 +from config import N_ITER + +# ========================================== +# 6. 网络维度配置导入 +# ========================================== +# 从全局配置导入网络维度设置 +from config import ( + NETWORK_USE_CUSTOM_DIMS, + NETWORK_BASE_DIM, + NETWORK_CUSTOM_DIMS, + NETWORK_POOL_RATIOS +) + +# ========================================== +# 0. BatchNorm 辅助函数 +# ========================================== +def get_batch_norm(num_features): + return nn.BatchNorm1d(num_features) + +# ========================================== +# 1. 物理感知图卷积层 +# ========================================== +class EMFullComplexLayer(MessagePassing): + def __init__(self, in_feats, out_feats): + super(EMFullComplexLayer, self).__init__() + self.aggr = 'mean' + total_in_dim = in_feats + in_feats + self.lin_fusion = nn.Linear(total_in_dim, out_feats) + + def forward(self, x, edge_index): + return self.propagate(edge_index, x=x) + + def message(self, x_j): + return x_j + + def update(self, aggs, x): + combined = torch.cat([x, aggs], dim=1) + fused_msg = self.lin_fusion(combined) + return fused_msg + +# ========================================== +# 2. 修改后的 GCN 模块 +# ========================================== +class GCN(nn.Module): + def __init__(self, in_feats, out_feats): + super(GCN, self).__init__() + self.conv = EMFullComplexLayer(in_feats, out_feats) + self.bn = get_batch_norm(out_feats) + self.gelu = nn.GELU() + + def forward(self, x, edge_index): + x = self.conv(x, edge_index) + x = self.bn(x) + x = self.gelu(x) + return x + +class FFTFeatureLayer(nn.Module): + """ + 对特征维度进行 FFT + """ + def __init__(self, in_feats, out_feats): + super(FFTFeatureLayer, self).__init__() + # 线性投影调整维度 + self.lin_in = nn.Linear(in_feats, out_feats) + + # 可学习的频域滤波器 (复数权重) + # RFFT 后，频域维度为 (out_feats // 2) + 1 + self.freq_dim = out_feats // 2 + 1 + + # 初始化复数权重 + # shape: [freq_dim] - 对每个频率分量进行缩放 + self.complex_weight = nn.Parameter( + torch.randn(self.freq_dim, 2, dtype=torch.float32) * 0.02 + ) + + def forward(self, x): + # x: [N, in_feats] -> [N, out_feats] + x = self.lin_in(x) + + # 修复：cuFFT在float16模式下只支持维度大小为2的幂次 + # 如果输入是float16且维度不是2的幂次，需要转换为float32 + original_dtype = x.dtype + converted_to_float32 = False + if x.dtype == torch.float16: + # 检查最后一个维度是否为2的幂次 + last_dim = x.size(-1) + is_power_of_two = (last_dim & (last_dim - 1)) == 0 and last_dim != 0 + if not is_power_of_two: + # 转换为float32以避免cuFFT限制 + x = x.to(torch.float32) + converted_to_float32 = True + + x_fft = torch.fft.rfft(x, dim=-1) + + # 2. 频域滤波 / 混合 + weight = torch.view_as_complex(self.complex_weight) + # 广播乘法: [N, freq_dim] * [freq_dim] + x_fft = x_fft * weight + # 3. 傅里叶逆变换 (Complex -> Real) + x_out = torch.fft.irfft(x_fft, n=x.size(-1), dim=-1) + + # 如果原始是float16且被转换过，转换回float16（保持精度一致性） + if original_dtype == torch.float16 and converted_to_float32: + x_out = x_out.to(original_dtype) + + return x_out + +# ========================================== +# 2.5 [核心修改] FFM 增强的输入层 +# ========================================== +class FFMEncodingLayer(nn.Module): + """ + 使用傅里叶特征映射 (FFM) 替换传统的线性输入层。 + 逻辑： + 1. 输入 features (如坐标或物理量) + 2. 映射 -> sin(2*pi*B*x), cos(2*pi*B*x) + 3. 线性投影融合 + 4. 物理图卷积 (EMFullComplexLayer) + """ + def __init__(self, in_feats, out_feats, sigma=1.0, mapping_size=None): + super(FFMEncodingLayer, self).__init__() + + # 1. 配置 FFM 参数 + self.input_dim = in_feats + # 如果未指定 mapping_size，默认设为 out_feats 的一半（因为sin/cos会翻倍） + # 或者设为一个固定的高维空间，例如 128 或 256 + self.mapping_size = mapping_size if mapping_size is not None else max(out_feats, 64) + self.sigma = sigma + + # 2. 初始化高斯随机矩阵 B (不可学习，类似位置编码) + # shape: [in_feats, mapping_size] + self.B = nn.Parameter( + torch.randn(self.input_dim, self.mapping_size) * self.sigma, + requires_grad=False + ) + + # FFM 后的维度 = mapping_size * 2 (sin + cos) + ffm_out_dim = self.mapping_size * 2 + + # 3. 特征融合层：将高维 FFM 特征投影回目标维度 + self.feature_projection = nn.Linear(ffm_out_dim, out_feats) + + # 4. 空间混合：保留图卷积能力 + self.conv = EMFullComplexLayer(out_feats, out_feats) + + # 5. 激活与归一化 + self.bn = get_batch_norm(out_feats) + self.gelu = nn.GELU() + + def input_mapping(self, x): + # x: [N, input_dim] + # 投影: (2 * pi * x) @ B + # 结果 shape: [N, mapping_size] + x_proj = (2.0 * torch.pi * x) @ self.B + + # 拼接 sin 和 cos -> [N, mapping_size * 2] + return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) + + def forward(self, x, edge_index): + # 1. FFM 映射 (提升频率感知能力) + x_ffm = self.input_mapping(x) + + # 2. 线性投影 (融合特征) + x_embed = self.feature_projection(x_ffm) + + # 3. 图卷积 (聚合邻居信息) + x_out = self.conv(x_embed, edge_index) + + # 4. 后处理 + x_out = self.bn(x_out) + x_out = self.gelu(x_out) + + return x_out + +# ========================================== +# 2.6 新增：混合谱图卷积层 (替代原 GCN) +# ========================================== +class SpectralGCN(nn.Module): + """ + 混合层：结合 局部图卷积(GCN) 和 频域特征变换(FFT) + Out = GCN(x, edge) + FFT(x) + """ + def __init__(self, in_feats, out_feats): + super(SpectralGCN, self).__init__() + + # 分支1：物理感知图卷积 (处理局部连接) + self.spatial_conv = EMFullComplexLayer(in_feats, out_feats) + + # 分支2：FFT 频域层 (处理全局/频域特征) + self.spectral_conv = FFTFeatureLayer(in_feats, out_feats) + + # 融合门控 (可学习的加权系数) + self.alpha = nn.Parameter(torch.tensor(0.5)) + + self.bn = get_batch_norm(out_feats) + self.gelu = nn.GELU() + + def forward(self, x, edge_index): + # 1. 空间路径：消息传递 + x_spatial = self.spatial_conv(x, edge_index) + + # 2. 频谱路径：FFT 变换 (不依赖 edge_index，纯节点特征变换) + x_spectral = self.spectral_conv(x) + + # 3. 残差融合 + # 使用 sigmoid 确保融合比例在 0-1 之间，或者直接相加 + # 这里使用加权求和，兼顾空间和频域信息 + x_out = x_spatial + self.alpha * x_spectral + + # 4. 激活与归一化 + x_out = self.bn(x_out) + x_out = self.gelu(x_out) + + return x_out + +# ========================================== +# 2.7 FFM处理 +# ========================================== +class ffm_process(nn.Module): + def __init__(self, in_feats, out_feats, ffm_sigma=1.0): + super(ffm_process, self).__init__() + self.B = nn.Parameter(torch.randn(in_feats, out_feats) * ffm_sigma, requires_grad=False) + + def forward(self, x): + x_proj = (2.0 * torch.pi * x) @ self.B + return torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) + +# ========================================== +# 3. Graph U-Nets 组件 (Checkpoint稳定版) +# ========================================== +class gPool(nn.Module): + """Graph Pooling layer - 简化版，不特殊处理背景场""" + def __init__(self, in_dim, ratio): + super(gPool, self).__init__() + self.ratio = ratio + self.proj = nn.Linear(in_dim, 1) # 评分投影 + self.sigmoid = nn.Sigmoid() + + def forward(self, x, edge_index): + N, C = x.size() + + # 计算节点重要性得分（使用所有特征） + scores = self.proj(x) # [N, 1] + scores = self.sigmoid(scores).squeeze() # [N] + + # 确定保留的节点数量 + k = max(1, int(self.ratio * N)) + + # 选择top-k节点 + values, idx = torch.topk(scores, k) + + # 门控机制：保留的节点特征按重要性加权 + x_pool = x[idx] * values.view(-1, 1) # [k, C] + + # 筛选保留节点之间的边 + row, col = edge_index + row_mask = torch.zeros(N, dtype=torch.bool, device=x.device) + col_mask = torch.zeros(N, dtype=torch.bool, device=x.device) + row_mask[idx] = True + col_mask[idx] = True + edge_mask = row_mask[row] & col_mask[col] + edge_index_pool = edge_index[:, edge_mask] + + # 重新映射节点索引到0-k范围 + idx_map = -torch.ones(N, dtype=torch.long, device=x.device) + idx_map[idx] = torch.arange(k, device=x.device) + edge_index_pool[0] = idx_map[edge_index_pool[0]] + edge_index_pool[1] = idx_map[edge_index_pool[1]] + + return x_pool, edge_index_pool, idx + +class gUnpool(nn.Module): + """Graph Unpooling layer (inverse of Pool)""" + def __init__(self): + super(gUnpool, self).__init__() + + def forward(self, x_pool, x_skip, idx): + # idx 对应 TopKPooling 返回的 perm (保留节点的索引) + N = x_skip.size(0) + C = x_pool.size(1) + + # 恢复原始图结构大小 + x_unpool = torch.zeros(N, C, device=x_pool.device, dtype=x_skip.dtype) + + if idx is not None and len(idx) > 0: + # 填充保留节点的特征 + # PyG TopKPooling 的 perm 保证了 idx < N + # 确保 x_pool 的行数与 idx 长度一致 + if x_pool.size(0) == len(idx): + x_unpool[idx] = x_pool + else: + # 理论上不应发生，除非维度对不上 + valid_len = min(x_pool.size(0), len(idx)) + x_unpool[idx[:valid_len]] = x_pool[:valid_len] + + # 跳跃连接（残差连接） + x_unpool = x_unpool + x_skip + + return x_unpool + +# ========================================== +# 4. 辅助模块 +# ========================================== +class StackedLinearBlock(nn.Module): + def __init__(self, in_feats, out_feats, dropout): + super(StackedLinearBlock, self).__init__() + self.fc = nn.Linear(in_feats, out_feats) + self.bn = get_batch_norm(out_feats) + self.gelu = nn.GELU() + self.dropout = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.fc.weight, a=1.0) + if self.fc.bias is not None: + nn.init.zeros_(self.fc.bias) + + def forward(self, x): + x = self.fc(x) + x = self.bn(x) + x = self.gelu(x) + x = self.dropout(x) + return x + +class FinalLinear(nn.Module): + def __init__(self, in_feats, out_feats): + super(FinalLinear, self).__init__() + self.fc = nn.Linear(in_feats, out_feats) + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.fc.weight, a=1.0) + if self.fc.bias is not None: + nn.init.zeros_(self.fc.bias) + + def forward(self, x): + x = self.fc(x) + return x + +# ========================================== +# 5. BuildGCN 网络构建 (最接近原始版本) +# ========================================== +class BuildGCN(nn.Module): + def __init__(self, input_feats, output_feats, base_dim=None, custom_dims=None): + super(BuildGCN, self).__init__() + + # 优先级：函数参数 > 全局配置 > 默认值 + if custom_dims is not None: + # 函数参数指定的自定义维度（最高优先级） + dims = custom_dims + elif base_dim is not None: + # 函数参数指定的基础维度 + dims = [base_dim, base_dim * 2, base_dim * 4] + elif NETWORK_USE_CUSTOM_DIMS: + # 全局配置的自定义维度 + dims = NETWORK_CUSTOM_DIMS + else: + # 全局配置的基础维度自动计算 + dims = [NETWORK_BASE_DIM, NETWORK_BASE_DIM * 2, NETWORK_BASE_DIM * 4] + + # 使用全局配置的池化比例 + pool_ratios = NETWORK_POOL_RATIOS + + # ========== 编码器路径 ========== + self.gcn1 = SpectralGCN(input_feats, dims[0]) + self.pool1 = TopKPooling(dims[0], pool_ratios[0]) + self.gcn2 = GCN(dims[0], dims[1]) + self.pool2 = TopKPooling(dims[1], pool_ratios[1]) + self.gcn3 = GCN(dims[1], dims[2]) + + # ========== 瓶颈层 ========== + self.bottom_gcn = GCN(dims[2], dims[2]) + + # ========== 解码器路径 ========== + self.dec_gcn1 = GCN(dims[2], dims[1]) + self.unpool1 = gUnpool() + self.dec_gcn2 = GCN(dims[1], dims[0]) + self.unpool2 = gUnpool() + self.dec_gcn3 = GCN(dims[0], dims[0]) + + # ========== 输出层 ========== + self.lin_stacked = StackedLinearBlock(dims[0], dims[0], dropout=0.5) + self.final_linear = FinalLinear(dims[0], output_feats) + + def forward(self, ndata, edata, batch=None): + x, edge_index = ndata, edata + + if batch is None: + batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device) + + # === 编码器路径 === + x1 = self.gcn1(x, edge_index) + + # 第1次池化 (TopKPooling) + # 返回值: x, edge_index, edge_attr, batch, perm, score + x2, edge_index2, _, batch2, idx1, _ = self.pool1(x1, edge_index, batch=batch) + + x2 = self.gcn2(x2, edge_index2) + + # 第2次池化 (TopKPooling) + x3, edge_index3, _, batch3, idx2, _ = self.pool2(x2, edge_index2, batch=batch2) + + x3 = self.gcn3(x3, edge_index3) + + # === 瓶颈层 === + x_bottom = self.bottom_gcn(x3, edge_index3) + + # === 解码器路径 === + # 注意：这里继续使用 edge_index3，和原逻辑一致 + x_up1 = self.dec_gcn1(x_bottom, edge_index3) + # 反池化：需要传入 idx (即 perm) 来恢复到上一层的大小 (x2 的大小) + x_up1 = self.unpool1(x_up1, x2, idx2) + + # 注意：这里使用 edge_index2 + x_up2 = self.dec_gcn2(x_up1, edge_index2) + # 反池化：恢复到输入层大小 (x1 的大小) + x_up2 = self.unpool2(x_up2, x1, idx1) + + x_out = self.dec_gcn3(x_up2, edge_index) + + # === 输出层 === + x_out = self.lin_stacked(x_out) + x_out = self.final_linear(x_out) + + return x_out + + +# ========================================== +# 7. BuildGCNList (保持不变) +# ========================================== +class BuildGCNList(nn.Module): + def __init__(self, input_feats, output_feats, n_iter=N_ITER, base_dim=None, custom_dims=None): + super(BuildGCNList, self).__init__() + self.n_iter = n_iter + self.input_feats = input_feats + self.output_feats = output_feats + self.networks = nn.ModuleList() + + for i in range(n_iter): + network = BuildGCN(input_feats, output_feats, base_dim=base_dim, custom_dims=custom_dims) + self.networks.append(network) + setattr(self, f'iter_{i}', network) + + print(f"已创建 {n_iter} 个集成Graph U-Nets的BuildGCN网络") + + def forward(self, ndata, edata, batch=None, iter_idx=None): + if iter_idx is not None: + if iter_idx < 0 or iter_idx >= self.n_iter: + raise ValueError(f"iter_idx必须在[0, {self.n_iter-1}]范围内") + return self.networks[iter_idx](ndata, edata, batch) + else: + outputs = [] + for i in range(self.n_iter): + output = self.networks[i](ndata, edata, batch) + outputs.append(output) + return outputs + + def get_network(self, iter_idx): + return self.networks[iter_idx] + + def __len__(self): + return self.n_iter \ No newline at end of file diff --git a/trainddp.py b/trainddp.py new file mode 100644 index 0000000..bf0c79d --- /dev/null +++ b/trainddp.py @@ -0,0 +1,2090 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau +from torch.utils.checkpoint import checkpoint # 梯度检查点，节省显存 +from datetime import timedelta # 用于DDP超时设置 +import matplotlib.pyplot as plt +# 混合精度训练已禁用，使用纯float32训练以保证稳定性 +from torch_geometric.loader import DataListLoader # 【关键】使用 ListLoader +import torch.distributed as dist # DDP 分布式训练 +import torch.multiprocessing as mp # 多进程 +from torch.nn.parallel import DistributedDataParallel as DDP # DDP +from torch_geometric.nn import DataParallel # 兼容性保留 +from torch_geometric.data import Batch, Data +import numpy as np +import os +import time +import copy +import json +from sklearn.model_selection import train_test_split + +# ========================================== +# 1. 导入自定义模块 +# ========================================== +from build_graph import build_graph_data, scan_all_data, root_data_path, load_file_data, is_main_process +from model import BuildGCNList +import os +# DDP环境优化：避免多进程数据加载导致的资源竞争 +# 在DDP中，每个GPU进程都会创建num_workers个子进程 +# 如果有N个GPU，num_workers=4，则总共有4*N个数据加载子进程 +# 这会导致文件描述符耗尽和共享内存对象过多 +# 建议：DDP环境中设置num_workers=0，让主进程处理数据加载 +NUM_WORKERS = 0 # DDP环境：禁用多进程数据加载以避免资源竞争 + +# ========================================== +# 2. 训练配置 +# ========================================== +# 学习率调度配置 +# ========================================== +# ReduceLROnPlateau 配置 +REDUCE_LR_PATIENCE = 30 # 多少个epoch无改善就降低学习率 +REDUCE_LR_FACTOR = 0.5 # 学习率降低因子 +REDUCE_LR_MIN_LR = 1e-6 # 最小学习率 +REDUCE_LR_MODE = 'min' # 'min'表示监控指标越小越好 +# REDUCE_LR_VERBOSE 已弃用，使用 get_last_lr() 替代 + +# ========================================== +# 早停机制配置 +# ========================================== +EARLY_STOPPING_PATIENCE = 50 # 连续多少个epoch无改善就停止训练 +EARLY_STOPPING_MIN_DELTA = 1e-5 # 最小改善阈值 (0.001e-2 = 1e-5) +EARLY_STOPPING_START_EPOCH = 10 # 从第几个epoch开始检查早停（给模型预热时间） +# DP 模式下 batch_size 是所有卡的总和 +# 例如：3张卡，batch_size=48 -> 每张卡分到 16 个图 +# 增加 batch size 以提高 GPU 利用率 +MATRIX_CACHE = {} +# 优化：DDP模式下每个GPU的batch size +# DDP性能优化：增大batch_size以更好地利用GPU并行计算 +# - 24GB 显存：建议 128-256 per GPU +# - 48GB 显存：建议 256-512 per GPU +# - 如果遇到OOM，可以适当减小 +TOTAL_BATCH_SIZE = 64 +EPOCH_ADAM = 2000 # 只使用Adam优化器 +TOTAL_EPOCHS = EPOCH_ADAM +EPOCH_PRINT = 50 +# 从全局配置文件导入参数 +from config import N_ITER, SAVE_DIR, OUTPUT_DIR, LOSS_TYPE, MASTER_PORT, LOAD_PRETRAINED_MODEL, PRETRAINED_MODEL_DIR + +# ========================================== +# 1.5. 输出目录配置 +# ========================================== +os.makedirs(OUTPUT_DIR, exist_ok=True) + +# 移除模块级别的打印，避免DDP重复打印 + +LR = 0.001 + +# 性能优化配置 +USE_AMP = False # 禁用混合精度训练（使用纯float32以保证稳定性） +# torch.compile 配置说明： +# - 需要 PyTorch 2.0+ 支持 +# - 可能提升 10-20% 前向传播速度 +# - 首次运行需要编译时间（可能较慢） +# - 可能与 DataParallel 和梯度检查点有兼容性问题 +# - 建议：先测试是否正常工作，如果遇到错误可以禁用 +USE_COMPILE = False # 启用torch.compile可提升10-20%速度（需要PyTorch 2.0+） +PIN_MEMORY = True # 启用pin_memory以加速数据传输 +USE_GRADIENT_CHECKPOINTING = False # 启用梯度检查点，节省显存（会稍微降低速度，约20-30%） + +# ========================================== +# 3. 全局矩阵缓存 (Multi-Device Matrix Cache) +# ========================================== +# 用于解决 DP 模式下不同显卡需要访问不同设备上矩阵的问题 + +def load_matrix_to_cache(data_mapping, n_total, device_ids, is_main_process=True): + """ + 将所有矩阵预加载到 CPU 上（使用 pin_memory 加速后续传输）。 + 结构: MATRIX_CACHE[k_idx] = (A_cpu, b_cpu) + + 优化：矩阵存储在 CPU，需要时才转移到 GPU，释放显存空间。 + 使用 pin_memory() 加速 CPU 到 GPU 的传输。 + + 注意: device_ids 参数保留用于兼容性，但不再为每个 GPU 创建副本 + """ + if is_main_process: + print(f"正在将矩阵预加载到 CPU（使用 pin_memory 优化）...") + print(f" 矩阵将在需要时动态传输到 GPU，以释放显存空间") + valid_indices = [k for k in range(1, n_total + 1) if k in data_mapping] + + for idx, k_idx in enumerate(valid_indices): + folder_path, folder_num, data_id = data_mapping[k_idx] + try: + # 读取原始数据 (CPU) + Aij = load_file_data(folder_path, "Aij", folder_num, data_id) + Av = load_file_data(folder_path, "Av", folder_num, data_id) + b_data = load_file_data(folder_path, "b", folder_num, data_id) + + rows = Aij[:, 0].astype(int) - 1 + cols = Aij[:, 1].astype(int) - 1 + values = Av[:, 0] + 1j * Av[:, 1] + N_nodes = len(b_data) + b_val = b_data[:, 0] + 1j * b_data[:, 1] + + shape = (N_nodes, N_nodes) + + # 在 CPU 上创建张量，并使用 pin_memory 加速后续传输 + i = torch.from_numpy(np.vstack((rows, cols))).long() + # v = torch.from_numpy(values.astype(np.complex128)) # 使用双精度complex + # b_k = torch.from_numpy(b_val.astype(np.complex128)) # 使用双精度complex + # 改为： + v = torch.from_numpy(values.astype(np.complex64)) # 32位复数 + b_k = torch.from_numpy(b_val.astype(np.complex64)) + + # # 将A矩阵和b向量都放大10倍 + # v = v * 10 + # b_k = b_k * 10 + + # 使用 pin_memory() 将数据固定在内存中，加速 CPU->GPU 传输 + # 注意：稀疏矩阵的 indices 和 values 可以 pin_memory + i = i.pin_memory() + v = v.pin_memory() + b_k = b_k.pin_memory() + + # 在 CPU 上创建稀疏矩阵（不 coalesce，延迟到 GPU 传输时） + # 存储 indices 和 values（已 pin_memory）以及 shape，而不是完整的稀疏矩阵 + # 这样可以保持 pin_memory 状态 + MATRIX_CACHE[k_idx] = { + 'indices': i, + 'values': v, + 'shape': shape, + 'b': b_k + } + + if is_main_process and (idx + 1) % 1000 == 0: + print(f" 已缓存 {idx + 1}/{len(valid_indices)} 个样本") + + except Exception as e: + if is_main_process: + print(f"加载样本 {k_idx} 出错: {e}") + continue + if is_main_process: + print("矩阵缓存完成（存储在 CPU，使用 pin_memory 优化，双精度complex128）。") + # 验证精度设置 + if MATRIX_CACHE: + sample_key = list(MATRIX_CACHE.keys())[0] + sample_data = MATRIX_CACHE[sample_key] + print(f"示例矩阵精度检查: values.dtype={sample_data['values'].dtype}, b.dtype={sample_data['b'].dtype}") + +def get_Ab(k_idx, device, dtype=None): + """ + 从缓存中获取当前设备对应的 A 和 b。 + 如果矩阵在 CPU 上，则动态传输到目标设备（使用 non_blocking 加速）。 + + Args: + k_idx: 样本索引 + device: 目标设备 + dtype: 目标数据类型，如果为None则保持原有精度 + + 优化： + 1. 矩阵存储在 CPU，需要时才传输到 GPU + 2. 使用 non_blocking=True 进行异步传输 + 3. 使用 pin_memory 加速传输 + """ + if k_idx not in MATRIX_CACHE: + raise RuntimeError(f"未找到样本 k={k_idx} 的缓存数据") + + cache_data = MATRIX_CACHE[k_idx] + + # 如果目标设备是 CPU，在 CPU 上构建稀疏矩阵 + if device.type == 'cpu': + indices = cache_data['indices'] + values = cache_data['values'] + shape = cache_data['shape'] + b_cpu = cache_data['b'] + + # 根据需要转换数据类型 + if dtype is not None: + values = values.to(dtype) + b_cpu = b_cpu.to(dtype) + + A_cpu = torch.sparse_coo_tensor(indices, values, shape, device=torch.device('cpu')) + return A_cpu, b_cpu + + # 如果目标设备是 GPU，将矩阵传输到 GPU + # 使用 non_blocking=True 进行异步传输（需要 pin_memory） + try: + # 获取已 pin_memory 的 indices 和 values + indices_cpu = cache_data['indices'] + values_cpu = cache_data['values'] + shape = cache_data['shape'] + b_cpu = cache_data['b'] + + # 根据需要转换数据类型 + if dtype is not None: + values_cpu = values_cpu.to(dtype) + b_cpu = b_cpu.to(dtype) + + # 异步传输到 GPU（non_blocking 需要 pin_memory） + indices_gpu = indices_cpu.to(device, non_blocking=True) + values_gpu = values_cpu.to(device, non_blocking=True) + + # 在 GPU 上重建稀疏矩阵并 coalesce + A_gpu = torch.sparse_coo_tensor(indices_gpu, values_gpu, shape, device=device).coalesce() + + # 异步传输 b 向量 + b_gpu = b_cpu.to(device, non_blocking=True) + + return A_gpu, b_gpu + except Exception as e: + # 如果异步传输失败，回退到同步传输 + if is_main_process(): + print(f"警告: 异步传输失败，使用同步传输 (k={k_idx}): {e}") + indices_cpu = cache_data['indices'] + values_cpu = cache_data['values'] + shape = cache_data['shape'] + b_cpu = cache_data['b'] + + # 根据需要转换数据类型 + if dtype is not None: + values_cpu = values_cpu.to(dtype) + b_cpu = b_cpu.to(dtype) + + indices_gpu = indices_cpu.to(device) + values_gpu = values_cpu.to(device) + A_gpu = torch.sparse_coo_tensor(indices_gpu, values_gpu, shape, device=device).coalesce() + b_gpu = b_cpu.to(device) + + return A_gpu, b_gpu + +# ========================================== +# 3. 辅助函数：绘制训练曲线 +# ========================================== +def plot_training_curve(train_losses=None, test_losses=None, save_path="training_curve.svg", data_file=None): + """ + 绘制训练集和测试集loss变化曲线并保存为文件 + + Args: + train_losses: 训练集loss列表（可选，如果提供data_file则忽略） + test_losses: 测试集loss列表（可选，如果提供data_file则忽略） + save_path: 保存路径 + data_file: 训练数据JSON文件路径，如果提供则从文件加载数据 + """ + # 在DDP环境中，只在主进程中执行绘图 + if not is_main_process(): + return + if data_file is not None: + # 从文件加载数据 + data = load_training_data(data_file) + if data is None: + if is_main_process(): + print("❌ 无法加载训练数据，跳过绘图") + return + train_losses = data.get('train_losses', []) + test_losses = data.get('test_losses', []) + if not train_losses or not test_losses: + if is_main_process(): + print("❌ 训练数据中缺少loss信息，跳过绘图") + return + plt.figure(figsize=(12, 6)) + + epochs = range(1, len(train_losses) + 1) + + # 绘制训练loss + plt.subplot(1, 2, 1) + plt.plot(epochs, train_losses, 'b-', label='Training Loss', linewidth=2) + plt.xlabel('Epoch') + plt.ylabel('Loss') + plt.title('Training Loss Curve') + plt.legend() + plt.grid(True, alpha=0.3) + plt.yscale('log') # 使用对数尺度 + + # 绘制测试loss + plt.subplot(1, 2, 2) + plt.plot(epochs, test_losses, 'r-', label='Test Loss', linewidth=2) + plt.xlabel('Epoch') + plt.ylabel('Loss') + plt.title('Test Loss Curve') + plt.legend() + plt.grid(True, alpha=0.3) + plt.yscale('log') # 使用对数尺度 + + # 保存为矢量图格式 + # PDF格式（高质量打印） + pdf_path = os.path.join(os.path.dirname(save_path), os.path.basename(save_path).replace('.png', '.pdf')) + plt.savefig(pdf_path, bbox_inches='tight') + plt.close() + + if is_main_process(): + print(f"✅ 训练曲线PDF已保存到: {pdf_path}") + + # SVG格式（网页和现代应用） + plt.figure(figsize=(12, 6)) + + # 重新绘制（合并在一个图中） + plt.plot(epochs, train_losses, 'b-', label='Training Loss', linewidth=2) + plt.plot(epochs, test_losses, 'r-', label='Test Loss', linewidth=2) + plt.xlabel('Epoch') + plt.ylabel('Loss') + plt.title('Training and Test Loss Curves') + plt.legend() + plt.grid(True, alpha=0.3) + plt.yscale('log') # 使用对数尺度 + + svg_path = os.path.join(os.path.dirname(save_path), os.path.basename(save_path).replace('.png', '.svg')) + plt.savefig(svg_path, bbox_inches='tight') + plt.close() + + if is_main_process(): + print(f"✅ 训练曲线SVG已保存到: {svg_path}") + + +def plot_mse_res_loss(train_mse=None, test_mse=None, train_res=None, test_res=None, save_path="mse_res_loss.svg", data_file=None): + """ + 绘制包含MSE loss和RES loss的训练曲线 + + Args: + train_mse: 训练集MSE loss列表（可选，如果提供data_file则忽略） + test_mse: 测试集MSE loss列表（可选，如果提供data_file则忽略） + train_res: 训练集RES loss列表（可选，如果提供data_file则忽略） + test_res: 测试集RES loss列表（可选，如果提供data_file则忽略） + save_path: 保存路径 + data_file: 训练数据JSON文件路径，如果提供则从文件加载数据 + """ + # 在DDP环境中，只在主进程中执行绘图 + if not is_main_process(): + return + if data_file is not None: + # 从文件加载数据 + data = load_training_data(data_file) + if data is None: + if is_main_process(): + print("❌ 无法加载训练数据，跳过绘图") + return + train_mse = data.get('train_mse_losses', []) + test_mse = data.get('test_mse_losses', []) + train_res = data.get('train_res_losses', []) + test_res = data.get('test_res_losses', []) + if not all([train_mse, test_mse, train_res, test_res]): + if is_main_process(): + print("❌ 训练数据中缺少MSE或RES loss信息，跳过绘图") + return + plt.figure(figsize=(14, 6)) + + epochs = range(1, len(train_mse) + 1) + + # 绘制MSE loss + plt.subplot(1, 2, 1) + plt.plot(epochs, train_mse, 'b-', label='Training MSE', linewidth=2, marker='o', markersize=3) + plt.plot(epochs, test_mse, 'r-', label='Testing MSE', linewidth=2, marker='+', markersize=3) + plt.xlabel('Epoch') + plt.ylabel('MSE Loss (log scale)') + plt.title('MSE Loss Curve') + plt.legend() + plt.grid(True, alpha=0.3) + plt.yscale('log') + + # 绘制RES loss + plt.subplot(1, 2, 2) + plt.plot(epochs, train_res, 'b-', label='Training Res', linewidth=2, marker='o', markersize=3) + plt.plot(epochs, test_res, 'r-', label='Testing Res', linewidth=2, marker='+', markersize=3) + plt.xlabel('Epoch') + plt.ylabel('RES loss') + plt.title('Residual Loss Curve') + plt.legend() + plt.grid(True, alpha=0.3) + plt.ylim(0, 0.2) # 设置y轴范围为0~0.2 + + # 保存为矢量图格式 (SVG) + svg_path = os.path.join(os.path.dirname(save_path), os.path.basename(save_path).replace('.png', '.svg')) + plt.savefig(svg_path, bbox_inches='tight') + plt.close() + + if is_main_process(): + print(f"✅ MSE和RES loss曲线SVG已保存到: {svg_path}") + + +def plot_mse_loss_distribution(train_mse_losses=None, test_mse_losses=None, save_path="mse_loss_distribution.svg", data_file=None): + """ + 绘制MSE loss分布柱状图 + + Args: + train_mse_losses: 训练集MSE loss列表（可选，如果提供data_file则忽略） + test_mse_losses: 测试集MSE loss列表（可选，如果提供data_file则忽略） + save_path: 保存路径 + data_file: 训练数据JSON文件路径，如果提供则从文件加载数据 + """ + # 在DDP环境中，只在主进程中执行绘图 + if not is_main_process(): + return + if data_file is not None: + # 从文件加载数据 + data = load_training_data(data_file) + if data is None: + if is_main_process(): + print("❌ 无法加载训练数据，跳过绘图") + return + train_mse_losses = data.get('train_mse_losses', []) + test_mse_losses = data.get('test_mse_losses', []) + if not train_mse_losses or not test_mse_losses: + if is_main_process(): + print("❌ 训练数据中缺少MSE loss信息，跳过绘图") + return + plt.figure(figsize=(12, 6)) + + # 确保输入是有效的数值列表 + if not train_mse_losses or not test_mse_losses: + if is_main_process(): + print("⚠️ 警告：训练集或测试集MSE loss数据为空，跳过绘图") + plt.close() + return + + # 转换为numpy数组并过滤无效值 + train_mse_losses = np.array(train_mse_losses, dtype=np.float32) + test_mse_losses = np.array(test_mse_losses, dtype=np.float32) + + # 过滤掉NaN和inf值 + train_mse_losses = train_mse_losses[np.isfinite(train_mse_losses)] + test_mse_losses = test_mse_losses[np.isfinite(test_mse_losses)] + + if len(train_mse_losses) == 0 or len(test_mse_losses) == 0: + if is_main_process(): + print("⚠️ 警告：过滤后训练集或测试集MSE loss数据为空，跳过绘图") + plt.close() + return + + # 计算统计信息 + train_mean = np.mean(train_mse_losses) + train_std = np.std(train_mse_losses) + test_mean = np.mean(test_mse_losses) + test_std = np.std(test_mse_losses) + + # 设置bins（对数坐标需要特殊的处理） + all_mse_losses = np.concatenate([train_mse_losses, test_mse_losses]) + # 确保所有值都是正数（MSE loss应该都是正数） + all_mse_losses = all_mse_losses[all_mse_losses > 0] + if len(all_mse_losses) == 0: + if is_main_process(): + print("⚠️ 警告：所有MSE loss值都是非正数，跳过绘图") + plt.close() + return + + # 创建对数bins + log_min = np.log10(max(all_mse_losses.min(), 1e-10)) # 避免log(0) + log_max = np.log10(all_mse_losses.max()) + bins = np.logspace(log_min, log_max, 50) + + # 计算直方图数据 + train_hist, _ = np.histogram(train_mse_losses, bins=bins, density=True) + test_hist, _ = np.histogram(test_mse_losses, bins=bins, density=True) + + # 计算bin中心用于绘制 + bin_centers = (bins[:-1] + bins[1:]) / 2 + + # 设置bar宽度 + bar_width = np.diff(bins) * 0.8 # 每个bin的80%宽度 + + # 绘制训练集和测试集的柱状图（分开放置） + plt.bar(bin_centers - bar_width/4, train_hist, width=bar_width/2, alpha=0.8, color='coral', + label=f'Training (μ={train_mean:.3f}, σ={train_std:.5f})', edgecolor='black', linewidth=0.5) + + plt.bar(bin_centers + bar_width/4, test_hist, width=bar_width/2, alpha=0.8, color='cyan', + label=f'Testing (μ={test_mean:.3f}, σ={test_std:.5f})', edgecolor='black', linewidth=0.5) + + # 拟合高斯分布并绘制 + from scipy import stats + + # 训练集高斯拟合 + train_params = stats.norm.fit(train_mse_losses) + train_x = np.logspace(log_min, log_max, 100) # 在对数空间均匀分布的点用于显示 + train_pdf = stats.norm.pdf(train_x, *train_params) + plt.plot(train_x, train_pdf, 'r-', linewidth=2, label='Training Gaussian Fit') + + # 测试集高斯拟合 + test_params = stats.norm.fit(test_mse_losses) + test_x = np.logspace(log_min, log_max, 100) # 在对数空间均匀分布的点用于显示 + test_pdf = stats.norm.pdf(test_x, *test_params) + plt.plot(test_x, test_pdf, 'b-', linewidth=2, label='Testing Gaussian Fit') + + plt.xlabel('MSE Loss (log scale)') + plt.ylabel('Density') + plt.title('MSE Loss Distribution') + plt.legend() + plt.grid(True, alpha=0.3) + plt.xscale('log') # 设置x轴为对数坐标 + + # 保存为矢量图格式 (SVG) + svg_path = os.path.join(os.path.dirname(save_path), os.path.basename(save_path).replace('.png', '.svg')) + plt.savefig(svg_path, bbox_inches='tight') + plt.close() + + if is_main_process(): + print(f"✅ MSE loss分布图SVG已保存到: {svg_path}") + print(f" 训练集: 均值={train_mean:.6f}, 标准差={train_std:.6f}") + print(f" 测试集: 均值={test_mean:.6f}, 标准差={test_std:.6f}") + + +def compute_relative_errors(solver, data_loader, data_mapping, device, matrix_dtype): + """ + 计算数据集中每个样本的相对误差 + + Args: + solver: 训练好的模型 + data_loader: 数据加载器 + data_mapping: 数据映射字典 + device: 设备 + matrix_dtype: 矩阵数据类型 + + Returns: + relative_errors: 每个样本的相对误差列表 + """ + solver.eval() + relative_errors = [] + + with torch.no_grad(): + for data_list in data_loader: + # 获取batch + if not isinstance(data_list, list): + data_list = [data_list] + + processed_data_list = [] + for item in data_list: + if isinstance(item, Data): + if not hasattr(item, 'k_idx'): + item.k_idx = torch.tensor([0]) + elif not isinstance(item.k_idx, torch.Tensor): + item.k_idx = torch.tensor([item.k_idx] if not isinstance(item.k_idx, (list, tuple)) else item.k_idx) + processed_data_list.append(item) + elif isinstance(item, tuple) and len(item) > 0 and isinstance(item[0], Data): + processed_data_list.append(item[0]) + + try: + batch = Batch.from_data_list(processed_data_list) + except: + from torch_geometric.data.collate import collate + batch, _, _ = collate(Data, processed_data_list, increment=True, add_batch=True, follow_batch=[]) + + batch = batch.to(device) + k_all = batch.k_idx + node_batch = batch.batch + B = k_all.size(0) + + # 获取模型精度 - 兼容不同层类型 + raw_model = solver.module if isinstance(solver, (DataParallel, DDP)) else solver + if hasattr(raw_model.model_real.networks[0].gcn1, 'conv'): + is_double_precision = raw_model.model_real.networks[0].gcn1.conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(raw_model.model_real.networks[0].gcn1, 'spatial_conv'): + is_double_precision = raw_model.model_real.networks[0].gcn1.spatial_conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(raw_model.model_real.networks[0].gcn1, 'linear'): + is_double_precision = raw_model.model_real.networks[0].gcn1.linear.weight.dtype == torch.float64 + else: + is_double_precision = False + + if is_double_precision: + data_dtype = torch.float64 + else: + data_dtype = torch.float32 + + # 准备数据 + eps_feat = batch.x[:, 0:2].to(data_dtype) + current_E_real = batch.x[:, 4].to(data_dtype) + current_E_imag = batch.x[:, 5].to(data_dtype) + true_real = batch.y[:, 0].to(data_dtype) + true_imag = batch.y[:, 1].to(data_dtype) + + # 加载矩阵 + A_list = [] + b_list = [] + for b_idx in range(B): + k = int(k_all[b_idx].item()) + A, b = get_Ab(k, device, matrix_dtype) + A_list.append(A) + b_list.append(b) + + # 前向传播 + E_real_cur = current_E_real + E_imag_cur = current_E_imag + + for iter_idx in range(raw_model.n_iter): + # 计算残差 + r_real_list = [] + r_imag_list = [] + for b_idx in range(B): + mask = (node_batch == b_idx) + E_r = E_real_cur[mask] + E_i = E_imag_cur[mask] + A = A_list[b_idx] + b_vec = b_list[b_idx] + E_c = torch.complex(E_r, E_i) + Ax = torch.sparse.mm(A, E_c.unsqueeze(-1)).squeeze(-1) + r_c = b_vec - Ax + if data_dtype == torch.float64: + r_real_list.append(r_c.real.double()) + r_imag_list.append(r_c.imag.double()) + else: + r_real_list.append(r_c.real.float()) + r_imag_list.append(r_c.imag.float()) + del E_c, Ax, r_c + + r_real = torch.cat(r_real_list, dim=0) + r_imag = torch.cat(r_imag_list, dim=0) + del r_real_list, r_imag_list + + # GNN前向 + # 从batch.x中提取背景场 + bg_real = batch.x[:, 6].to(data_dtype) + bg_imag = batch.x[:, 7].to(data_dtype) + x_in = torch.cat([ + eps_feat, + r_real.view(-1, 1), + r_imag.view(-1, 1), + E_real_cur.view(-1, 1), # 当前电场实部 (随迭代更新) + E_imag_cur.view(-1, 1), # 当前电场虚部 (随迭代更新) + bg_real.view(-1, 1), # 背景场实部 (不随网络更新) + bg_imag.view(-1, 1) # 背景场虚部 (不随网络更新) + ], dim=1) + + delta_real = raw_model.model_real(x_in, batch.edge_index, batch.batch, iter_idx) + delta_imag = raw_model.model_imag(x_in, batch.edge_index, batch.batch, iter_idx) + + E_real_cur = E_real_cur + delta_real.view(-1) + E_imag_cur = E_imag_cur + delta_imag.view(-1) + del x_in, delta_real, delta_imag, r_real, r_imag + + # 计算每个样本的相对误差 + for b_idx in range(B): + mask = (node_batch == b_idx) + pred_real = E_real_cur[mask] + pred_imag = E_imag_cur[mask] + true_r = true_real[mask] + true_i = true_imag[mask] + + # 计算平方的相对误差: ||pred - true||_2^2 / ||true||_2^2（与公式一致） + # 分子：||pred - true||_2^2 + numerator = ((pred_real - true_r).pow(2).sum() + (pred_imag - true_i).pow(2).sum()) + # 分母：||true||_2^2 + denominator = (true_r.pow(2).sum() + true_i.pow(2).sum()) + + if denominator > 1e-10: # 避免除零 + # 平方的相对误差：||x^FEM - x^GCN||_2^2 / ||x^FEM||_2^2 + rel_error_squared = (numerator / denominator).item() + relative_errors.append(rel_error_squared) + + del A_list, b_list, E_real_cur, E_imag_cur + + return relative_errors + +def save_training_data(training_data, save_path="training_data.json"): + """ + 保存训练数据到JSON文件 + + Args: + training_data: 包含训练数据的字典 + save_path: 保存路径 + """ + # 在DDP环境中，只在主进程中保存数据 + if not is_main_process(): + return + import json + + # 将numpy类型转换为Python原生类型，以便JSON序列化 + def convert_to_serializable(obj): + if isinstance(obj, dict): + return {k: convert_to_serializable(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [convert_to_serializable(item) for item in obj] + elif isinstance(obj, (np.integer, np.floating)): + return obj.item() + elif isinstance(obj, (int, float, str, bool)) or obj is None: + return obj + else: + return str(obj) + + serializable_data = convert_to_serializable(training_data) + + with open(save_path, 'w') as f: + json.dump(serializable_data, f, indent=2) + + if is_main_process(): + print(f"✅ 训练数据已保存到: {save_path}") + +def load_training_data(load_path="training_data.json"): + """ + 从JSON文件加载训练数据 + + Args: + load_path: 加载路径 + + Returns: + dict: 训练数据字典 + """ + import json + + try: + with open(load_path, 'r') as f: + data = json.load(f) + if is_main_process(): + print(f"✅ 训练数据已从 {load_path} 加载") + return data + except FileNotFoundError: + if is_main_process(): + print(f"⚠️ 警告：找不到训练数据文件 {load_path}") + return None + except Exception as e: + if is_main_process(): + print(f"❌ 加载训练数据失败: {e}") + return None + +# ========================================== +# 3. 辅助函数：处理 DataParallel 的输出 +# ========================================== +def extract_loss_and_num_nodes(outputs): + """ + 从 DataParallel 的输出中提取 loss 和 num_nodes。 + + 注意：forward 方法返回单个 tensor： + - 普通模式：[loss_sum, res_sum, num_nodes, num_samples] (4个值) + - Hybrid模式：[loss_sum, res_sum, num_nodes, num_samples, mse_loss_sum, phi_loss_sum] (6个值) + + Args: + outputs: DataParallel 返回的结果 + - 如果是列表，每个元素是上述格式的 tensor + - 如果是单个 tensor，形状为 [4]/[6] 或 [N, 4]/[N, 6]（N 个 GPU） + + Returns: + tuple: (total_loss_sum, total_res_sum, total_num_nodes, total_num_samples, total_mse_sum, total_phi_sum) + - total_loss_sum: 所有 GPU 的 loss 总和 (tensor) + - total_res_sum: 所有 GPU 的 RES loss 总和 (tensor) + - total_num_nodes: 所有 GPU 的 num_nodes 总和 (int) + - total_num_samples: 所有 GPU 的 num_samples 总和 (int) + - total_mse_sum: 所有 GPU 的 MSE loss 总和 (tensor，仅在hybrid模式下有效) + - total_phi_sum: 所有 GPU 的 Phi loss 总和 (tensor，仅在hybrid模式下有效) + """ + # 检查 outputs 是否为空（只检查 None 和空列表，不检查张量） + if outputs is None: + return torch.tensor(0.0), torch.tensor(0.0), 0, 0, torch.tensor(0.0), torch.tensor(0.0) + if isinstance(outputs, list) and len(outputs) == 0: + return torch.tensor(0.0), torch.tensor(0.0), 0, 0, torch.tensor(0.0), torch.tensor(0.0) + + # 处理 DataParallel 的输出格式 + # DataParallel 可能返回： + # 1. 列表: [[loss1, res1, n1, s1, mse1, phi1], [loss2, res2, n2, s2, mse2, phi2], ...] (hybrid模式) + # 2. 列表: [[loss1, res1, n1, s1], [loss2, res2, n2, s2], ...] (普通模式) + # 3. 单个 tensor: [[loss1, res1, n1, s1, mse1, phi1], ...] 形状为 [N, 6] 或 [N, 4] + # 4. 单个 tensor: [loss, res, n, s, mse, phi] 形状为 [6] 或 [4]（单 GPU 情况） + + if isinstance(outputs, list): + # 列表格式：每个元素是 [loss, res, num_nodes, num_samples, mse?, phi?] 的 tensor + loss_list = [] + res_list = [] + num_nodes_list = [] + num_samples_list = [] + mse_list = [] + phi_list = [] + is_hybrid_mode = False + + for o in outputs: + if isinstance(o, torch.Tensor): + if o.dim() == 1: + if o.shape[0] == 4: + # 普通模式：形状为 [4] 的 tensor + loss_list.append(o[0]) + res_list.append(o[1]) + num_nodes_list.append(o[2].item()) + num_samples_list.append(o[3].item()) + mse_list.append(torch.tensor(0.0)) # 占位符 + phi_list.append(torch.tensor(0.0)) # 占位符 + elif o.shape[0] == 6: + # Hybrid模式：形状为 [6] 的 tensor + is_hybrid_mode = True + loss_list.append(o[0]) + res_list.append(o[1]) + num_nodes_list.append(o[2].item()) + num_samples_list.append(o[3].item()) + mse_list.append(o[4]) + phi_list.append(o[5]) + else: + raise ValueError(f"不支持的 tensor 形状: {o.shape}") + elif o.dim() == 2: + if o.shape[1] == 4: + # 普通模式：形状为 [N, 4] 的 tensor（多个 GPU 合并） + loss_list.append(o[:, 0].sum()) + res_list.append(o[:, 1].sum()) + num_nodes_list.append(o[:, 2].sum().item()) + num_samples_list.append(o[:, 3].sum().item()) + mse_list.append(torch.tensor(0.0)) # 占位符 + phi_list.append(torch.tensor(0.0)) # 占位符 + elif o.shape[1] == 6: + # Hybrid模式：形状为 [N, 6] 的 tensor（多个 GPU 合并） + is_hybrid_mode = True + loss_list.append(o[:, 0].sum()) + res_list.append(o[:, 1].sum()) + num_nodes_list.append(o[:, 2].sum().item()) + num_samples_list.append(o[:, 3].sum().item()) + mse_list.append(o[:, 4].sum()) + phi_list.append(o[:, 5].sum()) + else: + raise ValueError(f"不支持的 tensor 形状: {o.shape}") + else: + raise ValueError(f"不支持的 tensor 形状: {o.shape}") + else: + raise TypeError(f"不支持的输出类型: {type(o)}") + + total_loss = torch.stack(loss_list).sum() if loss_list else torch.tensor(0.0) + total_res = torch.stack(res_list).sum() if res_list else torch.tensor(0.0) + total_num_nodes = sum(num_nodes_list) + total_num_samples = sum(num_samples_list) + total_mse = torch.stack(mse_list).sum() if is_hybrid_mode and mse_list else torch.tensor(0.0) + total_phi = torch.stack(phi_list).sum() if is_hybrid_mode and phi_list else torch.tensor(0.0) + + elif isinstance(outputs, torch.Tensor): + # 单个 tensor 格式 + if outputs.dim() == 1: + if outputs.shape[0] == 4: + # 普通模式：形状为 [4]：单 GPU + total_loss = outputs[0] + total_res = outputs[1] + total_num_nodes = int(outputs[2].item()) + total_num_samples = int(outputs[3].item()) + total_mse = torch.tensor(0.0) + total_phi = torch.tensor(0.0) + elif outputs.shape[0] == 6: + # Hybrid模式：形状为 [6]：单 GPU + total_loss = outputs[0] + total_res = outputs[1] + total_num_nodes = int(outputs[2].item()) + total_num_samples = int(outputs[3].item()) + total_mse = outputs[4] + total_phi = outputs[5] + elif outputs.shape[0] % 4 == 0 and outputs.shape[0] % 6 != 0: + # 普通模式：形状为 [4*N]：多个 GPU 的输出被展平（例如 [12] = 3个GPU * 4） + # 需要重塑为 [N, 4] 格式 + n_gpus = outputs.shape[0] // 4 + outputs_reshaped = outputs.view(n_gpus, 4) + total_loss = outputs_reshaped[:, 0].sum() + total_res = outputs_reshaped[:, 1].sum() + total_num_nodes = int(outputs_reshaped[:, 2].sum().item()) + total_num_samples = int(outputs_reshaped[:, 3].sum().item()) + total_mse = torch.tensor(0.0) + total_phi = torch.tensor(0.0) + elif outputs.shape[0] % 6 == 0: + # Hybrid模式：形状为 [6*N]：多个 GPU 的输出被展平 + # 需要重塑为 [N, 6] 格式 + n_gpus = outputs.shape[0] // 6 + outputs_reshaped = outputs.view(n_gpus, 6) + total_loss = outputs_reshaped[:, 0].sum() + total_res = outputs_reshaped[:, 1].sum() + total_num_nodes = int(outputs_reshaped[:, 2].sum().item()) + total_num_samples = int(outputs_reshaped[:, 3].sum().item()) + total_mse = outputs_reshaped[:, 4].sum() + total_phi = outputs_reshaped[:, 5].sum() + else: + raise ValueError(f"不支持的 tensor 形状: {outputs.shape}（一维张量长度必须是4或6的倍数）") + elif outputs.dim() == 2: + if outputs.shape[1] == 4: + # 普通模式：形状为 [N, 4]：多个 GPU 合并 + total_loss = outputs[:, 0].sum() + total_res = outputs[:, 1].sum() + total_num_nodes = int(outputs[:, 2].sum().item()) + total_num_samples = int(outputs[:, 3].sum().item()) + total_mse = torch.tensor(0.0) + total_phi = torch.tensor(0.0) + elif outputs.shape[1] == 6: + # Hybrid模式：形状为 [N, 6]：多个 GPU 合并 + total_loss = outputs[:, 0].sum() + total_res = outputs[:, 1].sum() + total_num_nodes = int(outputs[:, 2].sum().item()) + total_num_samples = int(outputs[:, 3].sum().item()) + total_mse = outputs[:, 4].sum() + total_phi = outputs[:, 5].sum() + else: + raise ValueError(f"不支持的 tensor 形状: {outputs.shape}") + else: + raise ValueError(f"不支持的 tensor 形状: {outputs.shape}") + else: + raise TypeError(f"不支持的输出类型: {type(outputs)}") + + return total_loss, total_res, total_num_nodes, total_num_samples, total_mse, total_phi + + +# ========================================== +# 4. Loss计算函数 +# ========================================== +def compute_mse_loss(E_real_cur, E_imag_cur, batch_y, num_nodes=None): + """ + 计算传统的MSE损失：||pred - true||^2（返回总和，未平均） + + Args: + E_real_cur: 预测的实部电场 [N] + E_imag_cur: 预测的虚部电场 [N] + batch_y: 真实标签 [N, 2] + num_nodes: 节点总数（用于接口统一，不使用） + + Returns: + loss_sum: MSE损失总和（未平均） + """ + true_real = batch_y[:, 0] + true_imag = batch_y[:, 1] + + # 计算实部和虚部的平方误差总和 + term1 = (E_real_cur - true_real).pow(2).sum() # 实部误差平方和 + term2 = (E_imag_cur - true_imag).pow(2).sum() # 虚部误差平方和 + + # 返回总SSE（Sum of Squared Errors） + loss_sum = 0.5 * (term1 + term2) + + # 释放中间变量 + del term1, term2 + + return loss_sum + +def compute_phi_loss(E_real_cur, E_imag_cur, k_all, node_batch, B, device, matrix_dtype, num_nodes=None): + """ + 计算Phi损失：||A*x - b||^2（物理残差，不需要节点平均） + + Args: + E_real_cur: 预测的实部电场 [N] + E_imag_cur: 预测的虚部电场 [N] + k_all: 样本索引 [B] + node_batch: 节点到批次的映射 [N] + B: 批次大小 + device: 计算设备 + matrix_dtype: 矩阵数据类型 + num_nodes: 节点总数（Phi loss不使用，用于接口统一） + + Returns: + loss_sum: Phi损失总和（不进行节点平均） + """ + phi_losses = [] + for b_idx in range(B): + k = int(k_all[b_idx].item()) + A, b = get_Ab(k, device, matrix_dtype) + + mask = (node_batch == b_idx) + E_r = E_real_cur[mask] + E_i = E_imag_cur[mask] + + E_c = torch.complex(E_r, E_i) + Ax = torch.sparse.mm(A, E_c.unsqueeze(-1)).squeeze(-1) + r = b - Ax + + # 计算残差的L2范数平方: ||Ax - b||_2^2 + phi_loss = torch.norm(r, p=2).pow(2) + phi_losses.append(phi_loss) + + # 使用torch.stack和torch.sum保持梯度计算图 + loss_sum = torch.stack(phi_losses).sum() + + return loss_sum + +def compute_asinh_loss(E_real_cur, E_imag_cur, batch_y, num_nodes): + """ + 计算Asinh损失：sqrt(asinh(||pred - true||^2))（不需要节点平均） + + Args: + E_real_cur: 预测的实部电场 [N] + E_imag_cur: 预测的虚部电场 [N] + batch_y: 真实标签 [N, 2] + num_nodes: 节点总数（Asinh loss不使用，用于接口统一） + + Returns: + loss_value: Asinh损失（不进行节点平均） + """ + true_real = batch_y[:, 0] + true_imag = batch_y[:, 1] + + # 计算预测值与真实值的差 + diff_real = E_real_cur - true_real + diff_imag = E_imag_cur - true_imag + + # 分别计算实部和虚部的平均平方误差 + mse_real = diff_real.pow(2).mean() # 实部平均平方误差 + mse_imag = diff_imag.pow(2).mean() # 虚部平均平方误差 + + # 分别应用asinh函数 + asinh_real = torch.asinh(mse_real) + asinh_imag = torch.asinh(mse_imag) + + # 将实部和虚部的asinh结果相加，然后开方 + loss_sum = torch.sqrt(asinh_real + asinh_imag) + + return loss_sum + +def compute_hybrid_loss(E_real_cur, E_imag_cur, batch_y, k_all, node_batch, B, device, matrix_dtype, epoch, num_nodes): + """ + 计算Hybrid损失：始终采用MSE + 0.1*Phi的固定组合（Phi loss进行节点平均） + + Args: + E_real_cur: 预测的实部电场 [N] + E_imag_cur: 预测的虚部电场 [N] + batch_y: 真实标签 [N, 2] + k_all: 样本索引 [B] + node_batch: 节点到批次的映射 [N] + B: 批次大小 + device: 计算设备 + matrix_dtype: 矩阵数据类型 + epoch: 当前训练轮次 + num_nodes: 节点总数 + + Returns: + tuple: (loss_value, mse_loss_sum, phi_loss_sum) + - loss_value: Hybrid损失（MSE + 0.1*Phi，进行Phi节点平均） + - mse_loss_sum: MSE loss总和（未平均） + - phi_loss_sum: Phi loss总和（未平均） + """ + # 计算MSE loss（总和，未平均） + mse_loss_sum = compute_mse_loss(E_real_cur, E_imag_cur, batch_y, num_nodes) + + # 计算Phi loss（总和，未平均） + phi_loss_sum = compute_phi_loss(E_real_cur, E_imag_cur, k_all, node_batch, B, device, matrix_dtype, num_nodes) + + # Phi loss进行节点平均：总Phi loss除以节点数num_nodes + phi_loss_per_node = phi_loss_sum / num_nodes + + # 固定权重策略：始终使用MSE + 0.1*Phi + lambda_phi = 0.5 + + # 组合损失：(MSE_sum / num_nodes) + 0.1 * Phi_per_node + loss_value = (mse_loss_sum / num_nodes) + lambda_phi * phi_loss_per_node + + return loss_value, mse_loss_sum, phi_loss_sum + +# ========================================== +# ========================================== +# 3.5. 预训练模型加载函数 +# ========================================== +def load_pretrained_model(model_dir, solver, n_iter, device, is_main_process): + """ + 只加载预训练模型权重（不加载优化器、调度器等状态） + 用于迁移学习：在新数据集上使用预训练模型继续训练 + + Args: + model_dir: 模型权重文件目录 + solver: 模型 + n_iter: 迭代次数 + device: 设备 + is_main_process: 是否为主进程 + + Returns: + bool: 是否成功加载 + """ + if is_main_process: + print(f"📂 正在加载预训练模型权重: {model_dir}") + + # 获取原始模型（去除DDP包装） + raw_model = solver.module if isinstance(solver, (DataParallel, DDP)) else solver + + # 检查模型文件是否存在 + all_exist = True + for i in range(n_iter): + real_path = os.path.join(model_dir, f"real_iter_{i}.pth") + imag_path = os.path.join(model_dir, f"imag_iter_{i}.pth") + if not os.path.exists(real_path) or not os.path.exists(imag_path): + all_exist = False + if is_main_process: + print(f"⚠️ 模型文件不存在: {real_path} 或 {imag_path}") + break + + if not all_exist: + if is_main_process: + print(" 将从随机初始化开始训练") + return False + + # 加载模型权重 + try: + for i in range(n_iter): + real_path = os.path.join(model_dir, f"real_iter_{i}.pth") + imag_path = os.path.join(model_dir, f"imag_iter_{i}.pth") + + real_net = raw_model.model_real.get_network(i) + imag_net = raw_model.model_imag.get_network(i) + + real_net.load_state_dict(torch.load(real_path, map_location=device, weights_only=True)) + imag_net.load_state_dict(torch.load(imag_path, map_location=device, weights_only=True)) + + if is_main_process: + print(f"✅ 预训练模型权重加载成功!") + print(f" 已加载 {n_iter} 个迭代网络的权重") + print(f" ⚠️ 注意：优化器、学习率调度器等状态已重新初始化") + + return True + except Exception as e: + if is_main_process: + print(f"❌ 加载预训练模型失败: {e}") + print(" 将从随机初始化开始训练") + return False + + +# 4. 物理求解器封装 (核心逻辑) +# ========================================== +class PhiSAGESolver(nn.Module): + """ + 将物理迭代循环封装为 Module，以便 DataParallel 可以自动分发计算。 + """ + def __init__(self, input_feats, output_feats, n_iter=N_ITER): + super(PhiSAGESolver, self).__init__() + self.n_iter = n_iter + # 内部实例化两个模型 + self.model_real = BuildGCNList(input_feats, output_feats, n_iter) + self.model_imag = BuildGCNList(input_feats, output_feats, n_iter) + + def forward(self, data_list, epoch=None): + """ + DP 模式下，data_list 是一个列表（原本 Batch 的一部分）。 + 我们需要在当前 GPU 上将其 Collate 成一个 Batch，然后跑物理循环。 + + Args: + data_list: 数据列表 + epoch: 当前训练轮次，用于hybrid loss计算（可选） + """ + # 1. 确保 data_list 中的元素都是 Data 对象 + # DataParallel 可能会传递特殊格式的数据，需要处理 + if not isinstance(data_list, list): + data_list = [data_list] + + # 检查并转换数据格式 + processed_data_list = [] + for item in data_list: + if isinstance(item, Data): + # 确保 Data 对象有必要的属性，并且 k_idx 是 tensor + if not hasattr(item, 'k_idx'): + # 如果没有 k_idx，尝试从其他属性获取或设置默认值 + item.k_idx = torch.tensor([0]) + elif not isinstance(item.k_idx, torch.Tensor): + item.k_idx = torch.tensor([item.k_idx] if not isinstance(item.k_idx, (list, tuple)) else item.k_idx) + processed_data_list.append(item) + elif isinstance(item, tuple): + # 如果是 tuple，可能是 (data, ...) 格式，取第一个元素 + if len(item) > 0 and isinstance(item[0], Data): + processed_data_list.append(item[0]) + else: + raise TypeError(f"无法处理的数据格式: {type(item)}, 内容: {item}") + else: + # 尝试直接使用，如果失败会抛出异常 + processed_data_list.append(item) + + # 2. 在当前 GPU 上构建 Batch + # 注意：如果遇到 tupleBatch 错误，可能是 PyG 版本问题 + # 尝试使用 collate 函数作为备选方案 + try: + batch = Batch.from_data_list(processed_data_list) + except (AttributeError, TypeError) as e: + error_msg = str(e) + if 'stores_as' in error_msg or 'tupleBatch' in error_msg: + # 使用 collate 函数手动构建 Batch + from torch_geometric.data.collate import collate + try: + batch, slice_dict, inc_dict = collate( + Data, + processed_data_list, + increment=True, + add_batch=True, + follow_batch=[], + ) + except Exception as e2: + # 如果 collate 也失败，提供更详细的错误信息 + if is_main_process(): + print(f"❌ Batch.from_data_list 失败: {e}") + print(f"❌ collate 也失败: {e2}") + print(f" 数据列表长度: {len(processed_data_list)}") + print(f" 第一个元素类型: {type(processed_data_list[0]) if processed_data_list else 'None'}") + if processed_data_list: + print(f" 第一个元素的属性: {dir(processed_data_list[0])}") + raise e + else: + raise e + + # 确保 batch 在正确的设备上 + # 在多卡情况下，DataParallel 会自动处理设备分配，batch 已经在正确的设备上 + # 在单卡情况下，需要确保 batch 在模型所在的设备上 + device = batch.x.device + + # 安全地获取模型设备（避免在 DataParallel replica 中出错） + try: + # 尝试获取模型参数所在的设备 + model_device = next(self.parameters()).device + # 如果 batch 不在模型设备上，则移动 batch（主要用于单卡情况） + if device != model_device: + batch = batch.to(model_device) + device = model_device + except (StopIteration, RuntimeError): + # 在 DataParallel 的 replica 中，参数可能不可用 + # 此时 batch 已经在正确的设备上（由 DataParallel 保证），直接使用 batch 的设备 + pass + + # 2. 准备数据 + k_all = batch.k_idx + node_batch = batch.batch + B = k_all.size(0) + + # 优化：减少不必要的 clone，使用 view 或直接索引 + # 节点特征：[eps_re, eps_im, r_re, r_im, Ebz_re, Ebz_im, bg_re, bg_im] + # 根据模型精度决定数据类型 - 兼容不同层类型 + if hasattr(self.model_real.networks[0].gcn1, 'conv'): + # GCN层的情况 + is_double = self.model_real.networks[0].gcn1.conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(self.model_real.networks[0].gcn1, 'spatial_conv'): + # SpectralGCN层的情况 + is_double = self.model_real.networks[0].gcn1.spatial_conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(self.model_real.networks[0].gcn1, 'linear'): + # FFTLayer的情况 + is_double = self.model_real.networks[0].gcn1.linear.weight.dtype == torch.float64 + else: + # 默认情况 + is_double = False + + if is_double: + eps_feat = batch.x[:, 0:2].double() # L-BFGS阶段用double + bg_real = batch.x[:, 6].double() # 背景场实部 (不随网络更新) + bg_imag = batch.x[:, 7].double() # 背景场虚部 (不随网络更新) + current_E_real = batch.x[:, 4].double() # 初始电场实部 + current_E_imag = batch.x[:, 5].double() # 初始电场虚部 + else: + eps_feat = batch.x[:, 0:2].float() # Adam阶段用float + bg_real = batch.x[:, 6].float() # 背景场实部 (不随网络更新) + bg_imag = batch.x[:, 7].float() # 背景场虚部 (不随网络更新) + current_E_real = batch.x[:, 4].float() # 初始电场实部 + current_E_imag = batch.x[:, 5].float() # 初始电场虚部 + + # 3. 优化：在物理迭代循环之前一次性加载所有矩阵到 GPU + # 这样可以在多次迭代中重复使用，避免重复的 CPU->GPU 传输（性能关键优化） + A_list = [] + b_list = [] + # 根据模型精度决定矩阵精度 - 兼容不同层类型 + if hasattr(self.model_real.networks[0].gcn1, 'conv'): + model_dtype = self.model_real.networks[0].gcn1.conv.lin_fusion.weight.dtype + elif hasattr(self.model_real.networks[0].gcn1, 'spatial_conv'): + model_dtype = self.model_real.networks[0].gcn1.spatial_conv.lin_fusion.weight.dtype + elif hasattr(self.model_real.networks[0].gcn1, 'linear'): + model_dtype = self.model_real.networks[0].gcn1.linear.weight.dtype + else: + model_dtype = torch.float32 + + matrix_dtype = torch.complex128 if model_dtype == torch.float64 else torch.complex64 + + for b_idx in range(B): + k = int(k_all[b_idx].item()) + A, b = get_Ab(k, device, matrix_dtype) + A_list.append(A) + b_list.append(b) + + # 只保存当前迭代的结果，而不是所有历史 + E_real_cur = current_E_real + E_imag_cur = current_E_imag + + # 4. 物理迭代循环 + for iter_idx in range(self.n_iter): + r_real_list = [] + r_imag_list = [] + + # 计算残差 r = b - A*E + # 优化：使用已加载的矩阵（已在 GPU 上），避免重复传输 + with torch.no_grad(): + for b_idx in range(B): + mask = (node_batch == b_idx) + E_r = E_real_cur[mask] + E_i = E_imag_cur[mask] + + # 使用已加载的矩阵（已在 GPU 上，无需重复传输） + A = A_list[b_idx] + b_vec = b_list[b_idx] + + E_c = torch.complex(E_r, E_i) + Ax = torch.sparse.mm(A, E_c.unsqueeze(-1)).squeeze(-1) + r_c = b_vec - Ax + + # 根据当前模型精度决定数据类型 - 兼容不同层类型 + if hasattr(self.model_real.networks[0].gcn1, 'conv'): + use_double = self.model_real.networks[0].gcn1.conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(self.model_real.networks[0].gcn1, 'spatial_conv'): + use_double = self.model_real.networks[0].gcn1.spatial_conv.lin_fusion.weight.dtype == torch.float64 + elif hasattr(self.model_real.networks[0].gcn1, 'linear'): + use_double = self.model_real.networks[0].gcn1.linear.weight.dtype == torch.float64 + else: + use_double = False + + if use_double: + r_real_list.append(r_c.real.double()) + r_imag_list.append(r_c.imag.double()) + else: + r_real_list.append(r_c.real.float()) + r_imag_list.append(r_c.imag.float()) + + # 释放中间变量（但保留矩阵 A 和 b，因为还要在下次迭代中使用） + del E_c, Ax, r_c + + r_real = torch.cat(r_real_list, dim=0) + r_imag = torch.cat(r_imag_list, dim=0) + + # 优化：释放中间列表 + del r_real_list, r_imag_list + + # 构造输入（优化：使用 view 而不是 unsqueeze，节省显存） + # 节点特征包含：[eps, r, E_current, bg] 共8个通道 + x_in = torch.cat([ + eps_feat, # eps_real, eps_imag [N, 2] + r_real.view(-1, 1), # r_real [N, 1] + r_imag.view(-1, 1), # r_imag [N, 1] + E_real_cur.view(-1, 1), # 当前电场实部 [N, 1] (随迭代更新) + E_imag_cur.view(-1, 1), # 当前电场虚部 [N, 1] (随迭代更新) + bg_real.view(-1, 1), # 背景场实部 [N, 1] (不随网络更新) + bg_imag.view(-1, 1) # 背景场虚部 [N, 1] (不随网络更新) + ], dim=1) + + # 优化：使用梯度检查点节省显存（在训练模式下） + use_checkpoint = USE_GRADIENT_CHECKPOINTING + + # 使用梯度检查点 + if self.training and use_checkpoint: + # 梯度检查点：在前向传播时不保存中间激活值，反向传播时重新计算 + # 这会节省显存，但会增加计算时间（约20-30%） + def gcn_forward_real(x, edge_index, batch, iter_idx): + return self.model_real(x, edge_index, batch, iter_idx) + def gcn_forward_imag(x, edge_index, batch, iter_idx): + return self.model_imag(x, edge_index, batch, iter_idx) + + delta_real = checkpoint(gcn_forward_real, x_in, batch.edge_index, batch.batch, iter_idx, use_reentrant=False) + delta_imag = checkpoint(gcn_forward_imag, x_in, batch.edge_index, batch.batch, iter_idx, use_reentrant=False) + else: + # 正常前向传播 + delta_real = self.model_real(x_in, batch.edge_index, batch.batch, iter_idx) + delta_imag = self.model_imag(x_in, batch.edge_index, batch.batch, iter_idx) + + # 展平并更新（使用 in-place 操作节省显存） + delta_real = delta_real.view(-1) + delta_imag = delta_imag.view(-1) + + # 优化：直接更新，不保存历史（只保留当前值） + E_real_next = E_real_cur + delta_real + E_imag_next = E_imag_cur + delta_imag + + # 优化：释放中间变量 + del x_in, delta_real, delta_imag, r_real, r_imag + + # 更新当前值（为下一次迭代准备） + E_real_cur = E_real_next + E_imag_cur = E_imag_next + + # 优化：不在每次迭代中清理显存，避免阻塞 + # 仅在阶段切换时清理显存碎片 + + # 5. 计算损失 + num_nodes = batch.x.size(0) # 总节点数 + if LOSS_TYPE == "hybrid": + # compute_hybrid_loss现在返回三个值，避免重复计算 + loss_sum, mse_loss_sum, phi_loss_sum = compute_hybrid_loss(E_real_cur, E_imag_cur, batch.y, k_all, node_batch, B, device, matrix_dtype, epoch, num_nodes) + else: + # 非hybrid模式下，只计算实际需要的loss + if LOSS_TYPE == "phi": + loss_sum = compute_phi_loss(E_real_cur, E_imag_cur, k_all, node_batch, B, device, matrix_dtype, num_nodes) + elif LOSS_TYPE == "asinh": + loss_sum = compute_asinh_loss(E_real_cur, E_imag_cur, batch.y, num_nodes) + else: + # 默认使用MSE loss + loss_sum = compute_mse_loss(E_real_cur, E_imag_cur, batch.y, num_nodes) + + # 非hybrid模式下不需要单独的MSE和Phi loss，使用占位符 + mse_loss_sum = torch.tensor(0.0, device=device) + phi_loss_sum = torch.tensor(0.0, device=device) + + # 6. 计算相对误差形式的RES loss + # RES loss需要真实标签，用于计算相对误差 + true_real = batch.y[:, 0] + true_imag = batch.y[:, 1] + # RES loss = ||x^FEM - x^GCN||_2^2 / ||x^FEM||_2^2 + # 其中 x^FEM 是真实解，x^GCN 是预测解 + # ||x||_2^2 = sum_i |x_i|^2，计算平方的相对误差（与公式一致） + res_loss_sum = torch.tensor(0.0, device=device, dtype=loss_sum.dtype) + with torch.no_grad(): + for b_idx in range(B): + mask = (node_batch == b_idx) + # 预测解 x^GCN + pred_real = E_real_cur[mask] + pred_imag = E_imag_cur[mask] + # 真实解 x^FEM + true_r = true_real[mask] + true_i = true_imag[mask] + + # 计算 ||x^FEM - x^GCN||_2^2 = sum_i |true_i - pred_i|^2 + # 对于复数向量，需要分别计算实部和虚部 + diff_real = true_r - pred_real + diff_imag = true_i - pred_imag + # 提前取绝对值再平方 + numerator = (torch.abs(diff_real).pow(2).sum() + torch.abs(diff_imag).pow(2).sum()) + + # 计算 ||x^FEM||_2^2 = sum_i |true_i|^2 + # 提前取绝对值再平方 + denominator = (torch.abs(true_r).pow(2).sum() + torch.abs(true_i).pow(2).sum()) + + # 避免除零，如果分母太小则使用一个小的epsilon + epsilon = 1e-10 + if denominator > epsilon: + # 平方的相对误差：||x^FEM - x^GCN||_2^2 / ||x^FEM||_2^2（与公式一致） + rel_error_squared = numerator / denominator + res_loss_sum = res_loss_sum + rel_error_squared + else: + # 如果真实解范数太小，使用绝对误差的平方 + res_loss_sum = res_loss_sum + numerator + + # 释放中间变量 + del diff_real, diff_imag, pred_real, pred_imag, true_r, true_i + + # 优化：在迭代结束后释放矩阵列表（释放显存） + del A_list, b_list + + # 释放变量 + del E_real_cur, E_imag_cur, true_real, true_imag + + # 注意：PyG DataParallel 可能不支持 tuple 返回值，会尝试将其当作 Batch 处理 + # 因此返回单个 tensor：[loss_sum, res_loss_sum, num_nodes, num_samples, mse_loss_sum, phi_loss_sum] + # 这样可以避免 'tupleBatch' 错误 + # num_nodes: 总节点数（用于loss的平均） + # num_samples: 总样本数（用于RES loss的平均，因为相对误差是针对每个样本的） + num_nodes = batch.x.size(0) # 总节点数（从 batch.x 获取，不需要从 E_real_cur） + num_samples = B # 样本数（batch中的图数量） + num_nodes_tensor = torch.tensor(num_nodes, dtype=torch.float32, device=device) + num_samples_tensor = torch.tensor(num_samples, dtype=torch.float32, device=device) + mse_loss_tensor = mse_loss_sum.to(dtype=torch.float32, device=device) if hasattr(mse_loss_sum, 'to') else torch.tensor(float(mse_loss_sum), dtype=torch.float32, device=device) + phi_loss_tensor = phi_loss_sum.to(dtype=torch.float32, device=device) if hasattr(phi_loss_sum, 'to') else torch.tensor(float(phi_loss_sum), dtype=torch.float32, device=device) + # 返回形状为 [6] 的 tensor: [loss_sum, res_loss_sum, num_nodes, num_samples, mse_loss_sum, phi_loss_sum] + return torch.stack([loss_sum, res_loss_sum, num_nodes_tensor, num_samples_tensor, mse_loss_tensor, phi_loss_tensor]) + + + +# ========================================== +# 5. 主程序 +# ========================================== +def setup_ddp(rank, world_size): + """设置DDP环境""" + os.environ['MASTER_ADDR'] = 'localhost' + os.environ['MASTER_PORT'] = MASTER_PORT + + # 设置NCCL优化环境变量，提高分布式训练稳定性 + os.environ['NCCL_TIMEOUT'] = '1800000' # 30分钟超时 (毫秒) + os.environ['NCCL_IB_DISABLE'] = '1' # 禁用IB以提高兼容性 + os.environ['NCCL_SOCKET_IFNAME'] = 'lo' # 使用本地环回接口 + os.environ['NCCL_DEBUG'] = 'WARN' # 设置调试级别 + + # 初始化进程组，设置更长的超时时间（30分钟）以避免NCCL超时 + dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(minutes=30)) + + # 设置当前进程的GPU + torch.cuda.set_device(rank) + device = torch.device(f'cuda:{rank}') + + return device + +def cleanup_ddp(): + """清理DDP环境""" + dist.destroy_process_group() + +def main_worker(rank, world_size): + """DDP训练的主工作函数""" + device = setup_ddp(rank, world_size) + + # 只有主进程(rank 0)输出信息 + is_main_process = (rank == 0) + + # DDP环境中，每个进程只负责一个GPU + device_ids = [rank] + + if is_main_process: + print(f"🚀 启动 DDP 分布式训练") + print(f" 进程 {rank}/{world_size}") + print(f" GPU: {device}") + print(f" 每个GPU的BatchSize: {TOTAL_BATCH_SIZE}") + print(f" 总BatchSize: {TOTAL_BATCH_SIZE * world_size}") + print(f" NCCL超时时间: 30分钟") + print(f" 环境变量: NCCL_TIMEOUT=1800000ms") + + # 打印数据集配置信息（只在主进程中打印） + from config import DATASET_TYPE, DATA_ROOT_PATH, SCA_PREFIX + print(f"📊 全局数据集配置: {DATASET_TYPE}") + print(f" 数据根目录: {DATA_ROOT_PATH}") + print(f" 文件命名前缀: {SCA_PREFIX}") + + # 2. 数据准备 + data_mapping, n_total = scan_all_data(root_data_path) + + # 优化：使用多进程并行构建数据集（如果数据量大） + dataset = [] + + # 获取所有有效的 k 索引 + valid_k_list = [k for k in range(1, n_total + 1) if k in data_mapping] + + # 使用串行方式构建数据集（多进程会导致共享内存问题） + # 注意：多进程传递大量 Data 对象时，Python 的 multiprocessing 使用共享内存 + # 可能导致 "Too many open files" 错误，因此使用串行方式更稳定 + if is_main_process: + print("构建数据集...") + print(" 串行构建数据集（稳定可靠）...") + for idx, k in enumerate(valid_k_list): + try: + data = build_graph_data(k) + if data is not None: + dataset.append(data) + # 每 100 个样本显示一次进度 + if (idx + 1) % 1000 == 0 and is_main_process: + print(f" 进度: {idx + 1}/{len(valid_k_list)} ({100*(idx+1)/len(valid_k_list):.1f}%)") + except Exception as e: + if (idx + 1) % 1000 == 0 and is_main_process: # 只在显示进度时打印错误 + print(f" 构建样本 {k} 失败: {e}") + pass + + if len(dataset) == 0: + raise RuntimeError("数据集为空，无法进行训练！") + + if is_main_process: + print(f" 数据集大小: {len(dataset)}") + train_ds, test_ds = train_test_split(dataset, test_size=0.2, random_state=42) + + # # 定义测试集的编号（从1开始计数，对应数据索引） + # test_indices = [1,13,14,18,19,24,25,36,37,42,43,48,49,60,61,66,67,72,73,84,85,90,91,96] + + # # 根据编号划分训练集和测试集 + # train_ds = [] + # test_ds = [] + + # for i, data in enumerate(dataset): + # # data.k_idx 存储的是数据的编号（从1开始） + # data_idx = data.k_idx.item() + # if data_idx in test_indices: + # test_ds.append(data) + # else: + # train_ds.append(data) + + if is_main_process: + print(f" 训练集: {len(train_ds)}, 测试集: {len(test_ds)}") + # print(f" 测试集编号: {test_indices}") + + # 【关键】使用 DataListLoader + # 注意：PyG 的 DataListLoader 可能不支持 num_workers（专为 DataParallel 设计） + # 但我们可以尝试设置，如果不支持会自动忽略 + # 优化：添加 pin_memory 以加速数据传输到GPU + + # Adam阶段：开启shuffle增加训练随机性，提高泛化能力 + train_loader = DataListLoader( + train_ds, + batch_size=TOTAL_BATCH_SIZE, + shuffle=True, # Adam阶段开启shuffle + drop_last=True, + pin_memory=True, # 确保为True + num_workers=NUM_WORKERS, # 添加这行！关键优化 + persistent_workers=True if NUM_WORKERS > 0 else False # 保持worker进程 +) + + # 测试阶段：关闭shuffle，确保评估结果一致性 + test_loader = DataListLoader( + test_ds, + batch_size=TOTAL_BATCH_SIZE, + shuffle=False, # 测试阶段关闭shuffle + pin_memory=True, + num_workers=NUM_WORKERS, # 添加 + persistent_workers=True if NUM_WORKERS > 0 else False +) + + # 3. 预加载矩阵到所有 GPU + load_matrix_to_cache(data_mapping, n_total, device_ids, is_main_process) + + # 【关键】DDP同步点：确保所有进程都完成矩阵预加载后再继续 + # 这可以避免进程间不同步导致的NCCL心跳超时 + if dist.is_initialized(): + dist.barrier() + if is_main_process: + print("✅ 所有进程已完成矩阵预加载，继续训练...") + + # 4. 模型初始化与 DP 包装 + # Adam阶段使用float32加快速度，L-BFGS阶段切换到float64提高精度 + # 节点特征包含：[eps, r, E_current, bg] 共8个通道 + solver = PhiSAGESolver(input_feats=8, output_feats=1, n_iter=N_ITER).float() # Adam阶段用float32 + solver.to(device) + + # 优化：使用 torch.compile 加速（如果支持） + # 注意：torch.compile 需要在 DataParallel 包装之前应用 + # 使用局部变量来避免修改全局变量 + use_compile = USE_COMPILE + if use_compile and is_main_process: + try: + # 检查 PyTorch 版本是否支持 compile + if hasattr(torch, 'compile'): + print("✅ 启用 torch.compile 优化...") + print(" ⚠️ 注意：首次运行需要编译时间，可能较慢") + print(" ⚠️ 如果遇到错误，请将 USE_COMPILE 设置为 False") + # 使用 'reduce-overhead' 模式，适合多次调用的场景 + solver = torch.compile(solver, mode='reduce-overhead') + else: + print("⚠️ PyTorch 版本不支持 torch.compile，跳过此优化") + use_compile = False + except Exception as e: + print(f"⚠️ torch.compile 失败: {e}，继续使用未编译版本") + print(f" 💡 提示：如果遇到兼容性问题，请将 USE_COMPILE 设置为 False") + use_compile = False + + # 打印网络维度信息 + if is_main_process: + from config import NETWORK_USE_CUSTOM_DIMS, NETWORK_BASE_DIM, NETWORK_CUSTOM_DIMS, NETWORK_POOL_RATIOS + + print("\n🔍 网络结构维度信息:") + print(f" 迭代次数 (n_iter): {solver.n_iter}") + print(f" 输入特征数: {solver.model_real.input_feats}") + print(f" 输出特征数: {solver.model_real.output_feats}") + + # 显示配置来源 + print(f" 配置来源: config.py") + if NETWORK_USE_CUSTOM_DIMS: + print(f" • 使用自定义维度: {NETWORK_CUSTOM_DIMS}") + else: + print(f" • 使用基础维度: {NETWORK_BASE_DIM} (自动计算: [{NETWORK_BASE_DIM}, {NETWORK_BASE_DIM*2}, {NETWORK_BASE_DIM*4}])") + print(f" • 池化配置: {NETWORK_POOL_RATIOS}") + + try: + # 获取第一个网络来显示维度信息 + first_network = solver.model_real.networks[0] + + # 显示网络维度配置 + if hasattr(first_network, 'gcn1') and hasattr(first_network.gcn1, 'conv') and hasattr(first_network.gcn1.conv, 'lin_fusion'): + gcn1_weight = first_network.gcn1.conv.lin_fusion.weight + gcn2_weight = first_network.gcn2.conv.lin_fusion.weight + gcn3_weight = first_network.gcn3.conv.lin_fusion.weight + + print(f" GCN层实际维度:") + print(f" • gcn1: {gcn1_weight.shape[1]} → {gcn1_weight.shape[0]}") + print(f" • gcn2: {gcn2_weight.shape[1]} → {gcn2_weight.shape[0]}") + print(f" • gcn3: {gcn3_weight.shape[1]} → {gcn3_weight.shape[0]}") + print(f" 网络架构: U-Net风格 (编码器-解码器)") + + # 计算参数量 + total_params = sum(p.numel() for p in solver.parameters()) + real_params = sum(p.numel() for p in solver.model_real.parameters()) + imag_params = sum(p.numel() for p in solver.model_imag.parameters()) + + print(f" 参数统计:") + print(f" • 总参数量: {total_params:,} ({total_params/1e6:.2f}M)") + print(f" • Real网络: {real_params:,} 参数") + print(f" • Imag网络: {imag_params:,} 参数") + print(f" • 单迭代网络: {real_params // solver.n_iter:,} 参数") + else: + print(" ⚠️ 无法获取详细的网络维度信息") + + except Exception as e: + print(f" ⚠️ 获取网络维度信息时出错: {e}") + # 仍然显示基本参数量信息 + total_params = sum(p.numel() for p in solver.parameters()) + print(f" 总参数量: {total_params:,} ({total_params/1e6:.2f}M)") + + # DDP包装 + solver = DDP(solver, device_ids=[rank], output_device=rank) + if is_main_process: + print("✅ 模型已通过 DDP 包装") + + # 使用纯float32训练，无需GradScaler + + # 5. 优化器 + # 注意：DP 包装后，参数名会多出 .module 前缀，但不影响 optimizer 识别 + # DDP优化：学习率按GPU数量线性缩放（总batch_size增大） + ddp_lr = LR # 每个GPU的基础学习率乘以GPU数量 + optimizer_adam = optim.Adam(solver.parameters(), lr=ddp_lr) + + # 使用ReduceLROnPlateau调度器，基于验证损失自动调整学习率 + scheduler = ReduceLROnPlateau( + optimizer_adam, + mode=REDUCE_LR_MODE, + factor=REDUCE_LR_FACTOR, + patience=REDUCE_LR_PATIENCE, + min_lr=REDUCE_LR_MIN_LR + ) + + # 6. 训练循环 + # 初始化训练状态变量 + best_loss = float('inf') # 用于早停判断的最佳loss + best_epoch = -1 + best_saved_loss = float('inf') # 用于保存模型判断的最佳loss（hybrid模式下200epoch后开始） + + # 早停机制变量 + early_stopping_counter = 0 + early_stopping_best_loss = float('inf') + + # 用于记录loss变化曲线 + train_losses = [] + test_losses = [] + train_mse_losses = [] + test_mse_losses = [] + train_res_losses = [] + test_res_losses = [] + + # 训练数据保存 + training_data = { + 'epochs': [], + 'train_losses': [], + 'test_losses': [], + 'train_mse_losses': [], + 'test_mse_losses': [], + 'train_res_losses': [], + 'test_res_losses': [], + 'train_relative_errors': [], + 'test_relative_errors': [] + } + + # Hybrid loss模式下的MSE和Phi loss记录 + if LOSS_TYPE == "hybrid": + hybrid_loss_data = { + 'epochs': [], + 'train_mse_losses': [], + 'train_phi_losses': [], + 'test_mse_losses': [], + 'test_phi_losses': [] + } + + # 记录训练总开始时间和上次打印时间 + total_start_time = time.time() + last_print_time = total_start_time + + # 加载预训练模型（如果启用） + if LOAD_PRETRAINED_MODEL: + pretrained_dir = PRETRAINED_MODEL_DIR if PRETRAINED_MODEL_DIR is not None else SAVE_DIR + load_success = load_pretrained_model(pretrained_dir, solver, N_ITER, device, is_main_process) + if load_success and is_main_process: + print("🔄 已加载预训练模型权重，优化器等状态已重新初始化") + print(" 将从epoch 0开始训练（在新数据集上）") + + for epoch in range(TOTAL_EPOCHS): + solver.train() + optimizer = optimizer_adam # 只使用Adam优化器 + + epoch_loss_sum = 0.0 + epoch_mse_sum = 0.0 + epoch_res_sum = 0.0 + epoch_phi_sum = 0.0 # Hybrid loss模式下的Phi loss累加器 + total_nodes = 0 + total_samples = 0 # 用于RES loss的平均（相对误差是针对每个样本的） + + # ========================== + # Adam Training (纯float32) + # ========================== + for data_list in train_loader: + optimizer.zero_grad() + + # Forward: list -> (split) -> GPUs -> (run) -> (gather) -> results + outputs = solver(data_list, epoch) + + # 使用辅助函数提取 loss 和 num_nodes + batch_loss_sum, batch_res_sum, num_nodes, num_samples, batch_mse_sum, batch_phi_sum = extract_loss_and_num_nodes(outputs) + + # 计算平均 Loss（Phi和Asinh loss不做节点平均，MSE loss使用节点平均） + # Hybrid loss根据当前权重决定是否节点平均 + if num_nodes == 0: + continue # 跳过空批次 + if LOSS_TYPE == "mse": + loss_mean = batch_loss_sum / num_nodes + else: + loss_mean = batch_loss_sum + + loss_mean.backward() + optimizer.step() + + epoch_loss_sum += batch_loss_sum.item() + epoch_res_sum += batch_res_sum.item() + # 只在hybrid模式下累加额外的loss分量 + if LOSS_TYPE == "hybrid": + epoch_mse_sum += batch_mse_sum.item() + epoch_phi_sum += batch_phi_sum.item() + total_nodes += num_nodes + total_samples += num_samples + + # 优化：每10个epoch清理一次显存 + if epoch % 10 == 0: + torch.cuda.empty_cache() + + + # ========================== + # 日志与测试 + # ========================== + num_batches = len(train_loader) + if num_batches > 0: + # 计算平均loss + if LOSS_TYPE == "mse": + # MSE loss: 除以总节点数，得到每个节点的平均loss + avg_train_loss = epoch_loss_sum / total_nodes if total_nodes > 0 else 0.0 + avg_train_mse = avg_train_loss # MSE模式下MSE loss就是总loss + elif LOSS_TYPE == "hybrid": + # Hybrid loss: 除以batch数，得到每个batch的平均loss + avg_train_loss = epoch_loss_sum / num_batches + avg_train_mse = epoch_mse_sum / total_nodes if total_nodes > 0 else 0.0 + else: + # Phi/Asinh loss: 除以batch数，得到每个batch的平均loss + avg_train_loss = epoch_loss_sum / num_batches + avg_train_mse = 0.0 # 非hybrid模式下不计算MSE loss + else: + avg_train_loss = 0.0 + avg_train_mse = 0.0 + + # RES Loss 是 Sum of Relative Errors，所以除以 total_samples (总图数) 是对的 + avg_train_res = epoch_res_sum / total_samples if total_samples > 0 else 0.0 + + # 计算平均Phi loss（在hybrid模式下） + if LOSS_TYPE == "hybrid": + avg_train_phi = epoch_phi_sum / total_nodes if total_nodes > 0 else 0.0 + else: + avg_train_phi = 0.0 + + # ========================== + # 测试集评估 (同理修正) + # ========================== + solver.eval() + test_loss_sum = 0.0 + test_mse_sum = 0.0 + test_res_sum = 0.0 + test_phi_sum = 0.0 # Hybrid loss模式下的Phi loss累加器 + test_total_nodes = 0 + test_total_samples = 0 + test_num_batches = len(test_loader) + + with torch.no_grad(): + # 评估阶段（纯float32） + for data_list in test_loader: + outputs = solver(data_list, epoch) + batch_loss_sum, batch_res_sum, num_nodes, num_samples, batch_mse_sum, batch_phi_sum = extract_loss_and_num_nodes(outputs) + + if num_nodes == 0: + continue + + # 累加基本loss + test_res_sum += batch_res_sum.item() # 累加RES loss + test_loss_sum += batch_loss_sum.item() * num_nodes # 转换为总损失再累加 + + # 只在hybrid模式下累加额外的loss分量 + if LOSS_TYPE == "hybrid": + test_mse_sum += batch_mse_sum.item() # 已经是总损失 + test_phi_sum += batch_phi_sum.item() # 已经是总损失 + test_total_nodes += num_nodes + test_total_samples += num_samples + test_num_batches += 1 + + if test_num_batches > 0: + # 按总节点数平均loss + if LOSS_TYPE == "mse": + # MSE loss: 除以总节点数，得到每个节点的平均loss + avg_test_loss = test_loss_sum / test_total_nodes if test_total_nodes > 0 else 0.0 + avg_test_mse = avg_test_loss # MSE模式下MSE loss就是总loss + elif LOSS_TYPE == "hybrid": + # Hybrid loss: 除以总节点数 + avg_test_loss = test_loss_sum / test_total_nodes if test_total_nodes > 0 else 0.0 + avg_test_mse = test_mse_sum / test_total_nodes if test_total_nodes > 0 else 0.0 + else: + # Phi/Asinh loss: 除以总节点数 + avg_test_loss = test_loss_sum / test_total_nodes if test_total_nodes > 0 else 0.0 + avg_test_mse = 0.0 # 非hybrid模式下不计算MSE loss + else: + avg_test_loss = 0.0 + avg_test_mse = 0.0 + + avg_test_res = test_res_sum / test_total_samples if test_total_samples > 0 else 0.0 + + # 计算平均Phi loss（在hybrid模式下） + if LOSS_TYPE == "hybrid": + avg_test_phi = test_phi_sum / test_total_nodes if test_total_nodes > 0 else 0.0 + else: + avg_test_phi = 0.0 + + # 使用ReduceLROnPlateau调度器 + scheduler.step(avg_test_loss) + + # Hybrid模式下从一开始就记录最佳loss并保存模型 + if LOSS_TYPE == "hybrid": + if avg_test_loss < best_saved_loss: + best_saved_loss = avg_test_loss + best_epoch = epoch + + # 保存模型 + raw_model = solver.module if isinstance(solver, (DataParallel, DDP)) else solver + os.makedirs(SAVE_DIR, exist_ok=True) + + # 保存模型权重 + n_iter = raw_model.n_iter + for i in range(n_iter): + torch.save(raw_model.model_real.get_network(i).state_dict(), + os.path.join(SAVE_DIR, f"real_iter_{i}.pth")) + torch.save(raw_model.model_imag.get_network(i).state_dict(), + os.path.join(SAVE_DIR, f"imag_iter_{i}.pth")) + else: + # 非hybrid模式，正常保存逻辑 + if avg_test_loss < best_saved_loss: + best_saved_loss = avg_test_loss + best_epoch = epoch + + # 保存模型 + raw_model = solver.module if isinstance(solver, (DataParallel, DDP)) else solver + os.makedirs(SAVE_DIR, exist_ok=True) + + # 保存模型权重 + n_iter = raw_model.n_iter + for i in range(n_iter): + torch.save(raw_model.model_real.get_network(i).state_dict(), + os.path.join(SAVE_DIR, f"real_iter_{i}.pth")) + torch.save(raw_model.model_imag.get_network(i).state_dict(), + os.path.join(SAVE_DIR, f"imag_iter_{i}.pth")) + + # 记录loss用于绘图 + train_losses.append(avg_train_loss) + test_losses.append(avg_test_loss) + train_mse_losses.append(avg_train_mse) + test_mse_losses.append(avg_test_mse) + train_res_losses.append(avg_train_res) + test_res_losses.append(avg_test_res) + + # 保存到训练数据字典 + training_data['epochs'].append(epoch) + training_data['train_losses'].append(avg_train_loss) + training_data['test_losses'].append(avg_test_loss) + training_data['train_mse_losses'].append(avg_train_mse) + training_data['test_mse_losses'].append(avg_test_mse) + training_data['train_res_losses'].append(avg_train_res) + training_data['test_res_losses'].append(avg_test_res) + + # 保存hybrid loss数据 + if LOSS_TYPE == "hybrid": + hybrid_loss_data['epochs'].append(epoch) + hybrid_loss_data['train_mse_losses'].append(avg_train_mse) + hybrid_loss_data['train_phi_losses'].append(avg_train_phi) + hybrid_loss_data['test_mse_losses'].append(avg_test_mse) + hybrid_loss_data['test_phi_losses'].append(avg_test_phi) + + # 全局早停检查（针对所有阶段） + if epoch >= EARLY_STOPPING_START_EPOCH: + early_stopping_enabled = True + if avg_test_loss < early_stopping_best_loss - EARLY_STOPPING_MIN_DELTA: + # 有显著改善，重置计数器 + early_stopping_best_loss = avg_test_loss + early_stopping_counter = 0 + else: + # 无显著改善，计数器加1 + early_stopping_counter += 1 + + # 检查是否达到早停条件 + if early_stopping_counter >= EARLY_STOPPING_PATIENCE: + if is_main_process: + print(f"🛑 全局早停激活！") + print(f" 连续{EARLY_STOPPING_PATIENCE}个epoch无显著改善") + print(f" 最小改善阈值: {EARLY_STOPPING_MIN_DELTA:.0e}") + print(f" 当前loss: {avg_test_loss:.6e}") + print(f" 最佳loss: {early_stopping_best_loss:.6e} (epoch {best_epoch})") + break + else: + # 预热阶段，跟踪最佳loss但不触发早停 + if avg_test_loss < early_stopping_best_loss: + early_stopping_best_loss = avg_test_loss + + + + if epoch % EPOCH_PRINT == 0 and is_main_process: + current_time = time.time() + interval_time = current_time - last_print_time + last_print_time = current_time + + # 为hybrid loss添加权重信息 + loss_info = f"Train Loss: {avg_train_loss:.6e} | Test Loss: {avg_test_loss:.6e}" + + print(f"Epoch {epoch:4d} | {loss_info} | " + f"Train RelErr: {avg_train_res:.6e} | Test RelErr: {avg_test_res:.6e} | " + f"Interval: {interval_time:.1f}s") + + # 训练结束，输出总耗时 + if is_main_process: + total_time = time.time() - total_start_time + hours = int(total_time // 3600) + minutes = int((total_time % 3600) // 60) + seconds = int(total_time % 60) + print(f"\n✅ 训练完成！总耗时: {hours:02d}:{minutes:02d}:{seconds:02d} ({total_time:.1f}秒)") + + # 输出最佳模型信息 + # 检查是否有模型被保存（通过best_epoch或best_saved_loss判断） + if best_epoch >= 0 and best_saved_loss < float('inf'): + print(f"\n🏆 最佳模型:") + print(f" 🎯 Epoch: {best_epoch}") + print(f" 📊 Test Loss: {best_saved_loss:.6e}") + print(f" 💾 模型已保存到: {SAVE_DIR}") + else: + print("⚠️ 警告：未找到有效的模型（可能训练失败）") + print(f" 📂 模型保存目录: {SAVE_DIR}") + + # 只在主进程中保存结果和生成图表 + if is_main_process: + # 保存训练数据 + training_data_path = os.path.join(OUTPUT_DIR, "training_data.json") + save_training_data(training_data, training_data_path) + + # 生成训练曲线 + curve_pdf_path = os.path.join(OUTPUT_DIR, "training_curve.pdf") + curve_svg_path = os.path.join(OUTPUT_DIR, "training_curve.svg") + plot_training_curve(train_losses, test_losses, curve_pdf_path) + + # 生成MSE和RES loss曲线 + mse_res_path = os.path.join(OUTPUT_DIR, "mse_res_loss.svg") + plot_mse_res_loss(train_mse_losses, test_mse_losses, train_res_losses, test_res_losses, mse_res_path) + + # 保存hybrid loss的MSE和Phi loss到txt文件（每50轮保存一次） + if LOSS_TYPE == "hybrid": + hybrid_loss_file = os.path.join(OUTPUT_DIR, "hybrid_loss_components.txt") + with open(hybrid_loss_file, 'w') as f: + f.write("Hybrid Loss Components (MSE + 1*Phi) - Every 50 epochs\n") + f.write("="*60 + "\n") + f.write("Epoch\tTrain_MSE\tTrain_Phi\tTest_MSE\tTest_Phi\n") + # 每50轮保存一次 + for i, epoch in enumerate(hybrid_loss_data['epochs']): + if epoch % 50 == 0: # 每50轮保存一次 + f.write(f"{epoch}\t{hybrid_loss_data['train_mse_losses'][i]:.6e}\t") + f.write(f"{hybrid_loss_data['train_phi_losses'][i]:.6e}\t") + f.write(f"{hybrid_loss_data['test_mse_losses'][i]:.6e}\t") + f.write(f"{hybrid_loss_data['test_phi_losses'][i]:.6e}\n") + + print(f"✅ Hybrid loss分量已保存到: {hybrid_loss_file}") + + # 计算相对误差分布 + if is_main_process: + print("\n正在计算相对误差分布...") + raw_model = solver.module if isinstance(solver, (DataParallel, DDP)) else solver + # 根据模型精度决定矩阵精度 - 兼容不同层类型 + if hasattr(raw_model.model_real.networks[0].gcn1, 'conv'): + model_precision = raw_model.model_real.networks[0].gcn1.conv.lin_fusion.weight.dtype + elif hasattr(raw_model.model_real.networks[0].gcn1, 'spatial_conv'): + model_precision = raw_model.model_real.networks[0].gcn1.spatial_conv.lin_fusion.weight.dtype + elif hasattr(raw_model.model_real.networks[0].gcn1, 'linear'): + model_precision = raw_model.model_real.networks[0].gcn1.linear.weight.dtype + else: + model_precision = torch.float32 + + matrix_dtype = torch.complex128 if model_precision == torch.float64 else torch.complex64 + + train_relative_errors = compute_relative_errors(solver, train_loader, data_mapping, device, matrix_dtype) + test_relative_errors = compute_relative_errors(solver, test_loader, data_mapping, device, matrix_dtype) + + # 保存相对误差数据 + training_data['train_relative_errors'] = train_relative_errors + training_data['test_relative_errors'] = test_relative_errors + + # 计算MSE loss分布（使用训练过程中的MSE loss值） + if is_main_process: + print("\n正在生成MSE loss分布图...") + train_mse_samples = [loss for loss in train_mse_losses for _ in range(10)] # 重复值以获得更好的分布 + test_mse_samples = [loss for loss in test_mse_losses for _ in range(10)] # 重复值以获得更好的分布 + + # 生成MSE loss分布图 + mse_dist_path = os.path.join(OUTPUT_DIR, "mse_loss_distribution.svg") + plot_mse_loss_distribution(train_mse_samples, test_mse_samples, mse_dist_path) + +def main(): + """主函数：启动DDP训练""" + # 检查是否有可用的GPU + if not torch.cuda.is_available(): + print("❌ 未检测到CUDA GPU，无法进行DDP训练") + return + + world_size = torch.cuda.device_count() + if world_size == 0: + print("❌ 未检测到任何GPU") + return + + print(f"🚀 启动DDP训练，使用 {world_size} 个GPU") + + # 使用spawn方式启动多进程 + try: + mp.spawn(main_worker, args=(world_size,), nprocs=world_size, join=True) + except KeyboardInterrupt: + print("\n⚠️ 训练被用户中断") + except Exception as e: + print(f"❌ DDP训练失败: {e}") + raise + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/visualize.py b/visualize.py new file mode 100644 index 0000000..e1468be --- /dev/null +++ b/visualize.py @@ -0,0 +1,281 @@ +import os +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.tri as mtri + +# 导入 build_graph 模块的函数和变量 +from build_graph import ( + scan_all_data, + root_data_path, + load_file_data, + data_mapping +) + +# 从全局配置导入数据集配置，确保文件命名格式同步 +from config import DATASET_TYPE, SCA_PREFIX + +# 注意：移除对save_test模块的导入以避免循环导入问题 +# k值现在通过函数参数传递，而不是模块级变量 + +# 默认算例编号 +DEFAULT_K = 8777 + +os.makedirs(root_data_path, exist_ok=True) + + +def load_data(k, use_prediction=True): + """ + 加载数据用于可视化（使用与build_graph.py一致的文件读取方式） + + Args: + k: 算例编号（全局索引） + use_prediction: 如果True，加载预测结果，否则加载真实值(Esz) + + Returns: + vertices: 节点坐标 [N, 2] + triangles: 三角形索引 [M, 3] + x_real: 实部 [N] + x_imag: 虚部 [N] + """ + # 扫描数据并获取data_mapping + data_mapping, n_total = scan_all_data(root_data_path) + + if k not in data_mapping: + raise ValueError(f"索引 k={k} 超出范围 (总数据量: {n_total})") + + folder_path, folder_num, data_id = data_mapping[k] + + # 读取节点坐标（vertex文件） + vertex_file = os.path.join(folder_path, f"vertex_{SCA_PREFIX}{folder_num}_{data_id}.txt") + if not os.path.exists(vertex_file): + raise FileNotFoundError(f"找不到节点坐标文件: {vertex_file}") + vertices = np.loadtxt(vertex_file) + + # 读取三角形索引（tri文件） + tri_file = os.path.join(folder_path, f"tri_{SCA_PREFIX}{folder_num}_{data_id}.txt") + if not os.path.exists(tri_file): + raise FileNotFoundError(f"找不到三角形索引文件: {tri_file}") + triangles = np.loadtxt(tri_file, dtype=int) - 1 # MATLAB索引转Python索引 + + # 根据use_prediction选择加载预测结果或真实值 + if use_prediction: + # 加载预测结果 + x_real_file = os.path.join(folder_path, f"Esz_pred_real_{SCA_PREFIX}{folder_num}_{data_id}.txt") + x_imag_file = os.path.join(folder_path, f"Esz_pred_imag_{SCA_PREFIX}{folder_num}_{data_id}.txt") + + if not os.path.exists(x_real_file) or not os.path.exists(x_imag_file): + raise FileNotFoundError(f"找不到预测结果文件: {x_real_file} 或 {x_imag_file}") + + x_real = np.loadtxt(x_real_file) + x_imag = np.loadtxt(x_imag_file) + else: + # 加载真实值（使用build_graph的load_file_data函数） + Esz_data = load_file_data(folder_path, "Esz", folder_num, data_id) + x_real = Esz_data[:, 0] + x_imag = Esz_data[:, 1] + + return vertices, triangles, x_real, x_imag + + +def load_mesh_data(k): + """ + 只加载网格数据（节点坐标和三角形索引），不加载场数据 + + Args: + k: 算例编号（全局索引） + + Returns: + vertices: 节点坐标 [N, 2] + triangles: 三角形索引 [M, 3] + """ + # 扫描数据并获取data_mapping + data_mapping, n_total = scan_all_data(root_data_path) + + if k not in data_mapping: + raise ValueError(f"索引 k={k} 超出范围 (总数据量: {n_total})") + + folder_path, folder_num, data_id = data_mapping[k] + + # 读取节点坐标（vertex文件） + vertex_file = os.path.join(folder_path, f"vertex_{SCA_PREFIX}{folder_num}_{data_id}.txt") + if not os.path.exists(vertex_file): + raise FileNotFoundError(f"找不到节点坐标文件: {vertex_file}") + vertices = np.loadtxt(vertex_file) + + # 读取三角形索引（tri文件） + tri_file = os.path.join(folder_path, f"tri_{SCA_PREFIX}{folder_num}_{data_id}.txt") + if not os.path.exists(tri_file): + raise FileNotFoundError(f"找不到三角形索引文件: {tri_file}") + triangles = np.loadtxt(tri_file, dtype=int) - 1 # MATLAB索引转Python索引 + + return vertices, triangles + + +def visualize_solution(k, output_dir="/public/home/zzx/gnn/PhiSAGE/PhiSAGE/visualizations", + use_prediction=True, save_combined=True, save_separate=False, + custom_data=None, custom_filename=None): + """ + 可视化解的实部、虚部和模值场图 + + Args: + k: 算例编号（全局索引） + output_dir: 输出目录 + use_prediction: 如果True，可视化预测结果，否则可视化真实值 + save_combined: 是否保存包含三个子图的组合图 + save_separate: 是否分别保存三个单独的图 + custom_data: 自定义数据字典，包含'vertices', 'triangles', 'x_real', 'x_imag'键 + custom_filename: 自定义文件名后缀，用于区分不同的迭代 + """ + if custom_data is not None: + # 使用提供的自定义数据 + vertices = custom_data['vertices'] + triangles = custom_data['triangles'] + x_real = custom_data['x_real'] + x_imag = custom_data['x_imag'] + else: + # 从文件加载数据 + vertices, triangles, x_real, x_imag = load_data(k, use_prediction=use_prediction) + + # 计算模值 + x_complex = x_real + 1j * x_imag + x_magnitude = np.abs(x_complex) + + # 创建输出目录 + os.makedirs(output_dir, exist_ok=True) + + x_coord = vertices[:, 0] + y_coord = vertices[:, 1] + triang = mtri.Triangulation(x_coord, y_coord, triangles) + + # 确定数据标签（用于文件名和标题） + if custom_filename is not None: + data_label = custom_filename + title_suffix = custom_filename + else: + data_label = "prediction" if use_prediction else "true" + title_suffix = f"k={k}" + + # 方案1: 保存组合图（三个子图） + if save_combined: + fig, axes = plt.subplots(1, 3, figsize=(18, 6)) + + # 实部 + tcf1 = axes[0].tricontourf(triang, x_real, levels=100, cmap='RdBu_r') + cbar1 = fig.colorbar(tcf1, ax=axes[0]) + cbar1.set_label('Re(u)', fontsize=12) + axes[0].set_title(f'Real Part ({title_suffix})', fontsize=14, fontweight='bold') + axes[0].set_aspect('equal') + axes[0].set_xlim(x_coord.min(), x_coord.max()) + axes[0].set_ylim(y_coord.min(), y_coord.max()) + axes[0].set_xticks([]) + axes[0].set_yticks([]) + + # 虚部 + tcf2 = axes[1].tricontourf(triang, x_imag, levels=100, cmap='RdBu_r') + cbar2 = fig.colorbar(tcf2, ax=axes[1]) + cbar2.set_label('Im(u)', fontsize=12) + axes[1].set_title(f'Imaginary Part ({title_suffix})', fontsize=14, fontweight='bold') + axes[1].set_aspect('equal') + axes[1].set_xlim(x_coord.min(), x_coord.max()) + axes[1].set_ylim(y_coord.min(), y_coord.max()) + axes[1].set_xticks([]) + axes[1].set_yticks([]) + + # 模值 + tcf3 = axes[2].tricontourf(triang, x_magnitude, levels=100, cmap='jet') + cbar3 = fig.colorbar(tcf3, ax=axes[2]) + cbar3.set_label('|u|', fontsize=12) + axes[2].set_title(f'Magnitude ({title_suffix})', fontsize=14, fontweight='bold') + axes[2].set_aspect('equal') + axes[2].set_xlim(x_coord.min(), x_coord.max()) + axes[2].set_ylim(y_coord.min(), y_coord.max()) + axes[2].set_xticks([]) + axes[2].set_yticks([]) + + plt.tight_layout() + out_file_combined = os.path.join(output_dir, f"{data_label}_combined.png") + plt.savefig(out_file_combined, dpi=300, bbox_inches='tight') + plt.close() + print(f"已保存组合图到: {out_file_combined}") + + # 方案2: 分别保存三个单独的图 + if save_separate: + # 实部 + fig, ax = plt.subplots(figsize=(8, 8)) + tcf = ax.tricontourf(triang, x_real, levels=100, cmap='RdBu_r') + cbar = fig.colorbar(tcf, ax=ax) + cbar.set_label('Re(u)', fontsize=14) + ax.set_title(f'Real Part ({title_suffix})', fontsize=16, fontweight='bold') + ax.set_aspect('equal') + ax.set_xlim(x_coord.min(), x_coord.max()) + ax.set_ylim(y_coord.min(), y_coord.max()) + ax.set_xticks([]) + ax.set_yticks([]) + plt.tight_layout() + out_file_real = os.path.join(output_dir, f"{data_label}_real.png") + plt.savefig(out_file_real, dpi=300, bbox_inches='tight') + plt.close() + print(f"已保存实部图到: {out_file_real}") + + # 虚部 + fig, ax = plt.subplots(figsize=(8, 8)) + tcf = ax.tricontourf(triang, x_imag, levels=100, cmap='RdBu_r') + cbar = fig.colorbar(tcf, ax=ax) + cbar.set_label('Im(u)', fontsize=14) + ax.set_title(f'Imaginary Part ({title_suffix})', fontsize=16, fontweight='bold') + ax.set_aspect('equal') + ax.set_xlim(x_coord.min(), x_coord.max()) + ax.set_ylim(y_coord.min(), y_coord.max()) + ax.set_xticks([]) + ax.set_yticks([]) + plt.tight_layout() + out_file_imag = os.path.join(output_dir, f"{data_label}_imag.png") + plt.savefig(out_file_imag, dpi=300, bbox_inches='tight') + plt.close() + print(f"已保存虚部图到: {out_file_imag}") + + # 模值 + fig, ax = plt.subplots(figsize=(8, 8)) + tcf = ax.tricontourf(triang, x_magnitude, levels=100, cmap='jet') + cbar = fig.colorbar(tcf, ax=ax) + cbar.set_label('|u|', fontsize=14) + ax.set_title(f'Magnitude ({title_suffix})', fontsize=16, fontweight='bold') + ax.set_aspect('equal') + ax.set_xlim(x_coord.min(), x_coord.max()) + ax.set_ylim(y_coord.min(), y_coord.max()) + ax.set_xticks([]) + ax.set_yticks([]) + plt.tight_layout() + out_file_mag = os.path.join(output_dir, f"{data_label}_magnitude.png") + plt.savefig(out_file_mag, dpi=300, bbox_inches='tight') + plt.close() + print(f"已保存模值图到: {out_file_mag}") + + +if __name__ == "__main__": + # 可视化指定算例 + # 尝试从 save_test 模块读取 k 值，避免循环导入 + try: + # 使用动态导入来避免启动时的循环导入问题 + import sys + if 'save_test' in sys.modules: + save_test_module = sys.modules['save_test'] + current_k = getattr(save_test_module, 'k', DEFAULT_K) + print(f"从已加载的 save_test 模块读取到 k={current_k}") + else: + current_k = DEFAULT_K + print(f"save_test 模块未加载，使用默认 k={current_k}") + except Exception as e: + current_k = DEFAULT_K + print(f"读取 k 值失败: {e}，使用默认 k={current_k}") + + # 可视化预测结果（默认） + print(f"正在可视化算例 k={current_k} 的预测结果...") + visualize_solution(k=current_k, use_prediction=True, save_combined=True, save_separate=False) + + # 可选：同时可视化真实值进行对比 + print(f"正在可视化算例 k={current_k} 的真实值...") + visualize_solution(k=current_k, use_prediction=False, save_combined=True, save_separate=False) + + +