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()