From 771709fc4b7850b394ac47bee8a8e010233ef326 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Nov 2025 14:13:36 +0000 Subject: [PATCH 1/3] Add setup_config.sh script for traffic data imputation configuration - Create interactive configuration wizard for LD2011_2014 dataset - Support GPU selection (A10/A40) and missing rate configuration (20%-80%) - Auto-generate training configs and launch scripts - Add comprehensive Chinese comments and usage instructions --- setup_config.sh | 217 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100755 setup_config.sh diff --git a/setup_config.sh b/setup_config.sh new file mode 100755 index 0000000..0b2f43f --- /dev/null +++ b/setup_config.sh @@ -0,0 +1,217 @@ +#!/bin/bash +# ============================================================================ +# 交通数据插补隐式显式扩散模型配置脚本 +# 用途：自动创建针对LD2011_2014.txt数据集的配置文件 +# ============================================================================ + +# 设置颜色输出 +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' # No Color + +echo -e "${GREEN}========================================${NC}" +echo -e "${GREEN}交通数据插补模型配置向导${NC}" +echo -e "${GREEN}========================================${NC}" + +# 检查是否在正确的目录 +if [ ! -d "src" ]; then + echo -e "${RED}错误: 请在项目根目录运行此脚本${NC}" + exit 1 +fi + +# 创建配置目录（如果不存在） +CONFIG_DIR="src/config" +mkdir -p "$CONFIG_DIR" + +# 询问用户GPU设置 +echo -e "\n${YELLOW}请选择要使用的GPU:${NC}" +echo "1) cuda:0 (A40)" +echo "2) cuda:1 (A10)" +read -p "请输入选项 (1 或 2): " gpu_choice + +case $gpu_choice in + 1) + GPU_ID=0 + GPU_NAME="A40" + ;; + 2) + GPU_ID=1 + GPU_NAME="A10" + ;; + *) + echo -e "${YELLOW}无效选择，默认使用 cuda:0 (A40)${NC}" + GPU_ID=0 + GPU_NAME="A40" + ;; +esac + +echo -e "${GREEN}✓ 已选择 GPU: cuda:${GPU_ID} (${GPU_NAME})${NC}" + +# 询问缺失率 +echo -e "\n${YELLOW}请选择要训练的缺失率:${NC}" +echo "1) 20%" +echo "2) 30%" +echo "3) 40%" +echo "4) 50%" +echo "5) 60%" +echo "6) 70%" +echo "7) 80%" +echo "8) 全部 (依次训练所有缺失率)" +read -p "请输入选项 (1-8): " missing_choice + +case $missing_choice in + 1) MISSING_RATES=(20) ;; + 2) MISSING_RATES=(30) ;; + 3) MISSING_RATES=(40) ;; + 4) MISSING_RATES=(50) ;; + 5) MISSING_RATES=(60) ;; + 6) MISSING_RATES=(70) ;; + 7) MISSING_RATES=(80) ;; + 8) MISSING_RATES=(20 30 40 50 60 70 80) ;; + *) + echo -e "${YELLOW}无效选择，默认使用 20%${NC}" + MISSING_RATES=(20) + ;; +esac + +echo -e "${GREEN}✓ 已选择缺失率: ${MISSING_RATES[@]}%${NC}" + +# 数据集路径 +DATA_PATH="/home/zhu/sssdtcn/LD2011_2014.txt" + +# 检查数据集是否存在 +if [ ! -f "$DATA_PATH" ]; then + echo -e "${RED}警告: 数据集文件不存在: $DATA_PATH${NC}" + echo -e "${YELLOW}请确保数据集已正确放置${NC}" +fi + +# 为每个缺失率创建配置文件 +for MISSING_RATE in "${MISSING_RATES[@]}"; do + echo -e "\n${YELLOW}正在创建缺失率 ${MISSING_RATE}% 的配置文件...${NC}" + + CONFIG_FILE="${CONFIG_DIR}/config_traffic_${MISSING_RATE}.json" + + cat > "$CONFIG_FILE" << EOF +{ + "diffusion_config": { + "T": 200, + "beta_0": 0.0001, + "beta_T": 0.02 + }, + "wavenet_config": { + "in_channels": 370, + "out_channels": 370, + "num_res_layers": 36, + "res_channels": 256, + "skip_channels": 256, + "diffusion_step_embed_dim_in": 128, + "diffusion_step_embed_dim_mid": 512, + "diffusion_step_embed_dim_out": 512, + "s4_lmax": 168, + "s4_d_state": 64, + "s4_dropout": 0.0, + "s4_bidirectional": 1, + "s4_layernorm": 1, + "tcn_channels": [256, 256, 256], + "tcn_kernel_size": 3, + "tcn_dilation_rates": [1, 2, 4, 8] + }, + "train_config": { + "output_directory": "./results/traffic/${MISSING_RATE}", + "ckpt_iter": "max", + "iters_per_ckpt": 1000, + "iters_per_logging": 100, + "n_iters": 50000, + "learning_rate": 2e-4, + "batch_size": 8, + "only_generate_missing": 1, + "use_model": 2, + "masking": "rm", + "missing_k": ${MISSING_RATE}, + "gpu_id": ${GPU_ID} + }, + "trainset_config": { + "train_data_path": "${DATA_PATH}", + "test_data_path": "${DATA_PATH}", + "segment_length": 168, + "sampling_rate": 168, + "train_split": 0.8 + }, + "gen_config": { + "output_directory": "./results/traffic/${MISSING_RATE}", + "ckpt_path": "./results/traffic/${MISSING_RATE}/" + } +} +EOF + + echo -e "${GREEN}✓ 已创建配置文件: $CONFIG_FILE${NC}" +done + +# 创建环境变量设置脚本 +ENV_FILE="set_gpu_env.sh" +cat > "$ENV_FILE" << EOF +#!/bin/bash +# GPU环境变量设置 +export CUDA_VISIBLE_DEVICES=${GPU_ID} +export CUDA_DEVICE_ORDER=PCI_BUS_ID +echo "已设置 CUDA_VISIBLE_DEVICES=${GPU_ID}" +EOF + +chmod +x "$ENV_FILE" +echo -e "\n${GREEN}✓ 已创建GPU环境设置脚本: $ENV_FILE${NC}" + +# 创建训练启动脚本 +TRAIN_SCRIPT="run_training.sh" +cat > "$TRAIN_SCRIPT" << EOF +#!/bin/bash +# 交通数据插补模型训练脚本 + +# 设置GPU环境 +source ./set_gpu_env.sh + +# 切换到src目录 +cd src + +# 训练模型 +for MISSING_RATE in ${MISSING_RATES[@]}; do + echo "==========================================" + echo "开始训练缺失率 \${MISSING_RATE}% 的模型" + echo "==========================================" + + # 清理之前的结果 + rm -rf ./results/traffic/\${MISSING_RATE}/* + + # 开始训练 + python3 train.py -c config/config_traffic_\${MISSING_RATE}.json + + # 运行推理 + echo "开始推理和评估..." + python3 inference.py -c config/config_traffic_\${MISSING_RATE}.json + + echo "缺失率 \${MISSING_RATE}% 训练完成！" + echo "" +done + +echo "所有训练任务完成！" +EOF + +chmod +x "$TRAIN_SCRIPT" +echo -e "${GREEN}✓ 已创建训练启动脚本: $TRAIN_SCRIPT${NC}" + +# 显示后续步骤 +echo -e "\n${GREEN}========================================${NC}" +echo -e "${GREEN}配置完成！${NC}" +echo -e "${GREEN}========================================${NC}" +echo -e "\n${YELLOW}后续步骤:${NC}" +echo -e "1. 安装依赖包:" +echo -e " ${GREEN}cd src && pip install -r requirements.txt${NC}" +echo -e "\n2. 开始训练:" +echo -e " ${GREEN}./run_training.sh${NC}" +echo -e "\n3. 查看结果:" +echo -e " ${GREEN}ls -la results/traffic/${NC}" +echo -e "\n${YELLOW}提示:${NC}" +echo -e "- 配置文件位置: src/config/config_traffic_XX.json" +echo -e "- GPU设置: cuda:${GPU_ID} (${GPU_NAME})" +echo -e "- 数据集: ${DATA_PATH}" +echo "" From 6f39a179725f85f327e966161e6dff26edd41047 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Nov 2025 14:33:30 +0000 Subject: [PATCH 2/3] Add complete IEclaude implementation for traffic data imputation This commit adds a complete Implicit-Explicit Diffusion Model for traffic data imputation: Features: - TCN Implicit Feature Extractor: Multi-scale temporal feature extraction with dilated causal convolutions - S4 Explicit Feature Extractor: Long-term dependency modeling with Structured State Space Models - IE Diffusion Model: Integration of implicit and explicit features for enhanced imputation - Complete training and evaluation pipeline - Support for multiple missing rates (20%-80%) - Comprehensive Chinese comments on every line of code Components: 1. models/: Core model implementations (TCN, S4, IEDiffusion) 2. data/: Traffic data loader and preprocessing 3. utils/: Diffusion process utilities 4. train.py: Training script with multi-GPU support 5. evaluate.py: Evaluation script with MAE/RMSE metrics 6. configs/: Configuration files for different missing rates 7. run_all.sh: One-click script to run all experiments 8. README.md: Detailed documentation with: - Architecture design and explanation - Hyperparameter tuning guide - Code explanation for each module - Quick start guide Documentation: - Every line of code has detailed Chinese comments - Three-level hyperparameter tuning guide (beginner, intermediate, advanced) - Explanations of TCN dilation rates, S4 ABCD matrices, and diffusion parameters --- IEclaude/README.md | 803 ++++++++++++++++++++++++++++ IEclaude/configs/config_20.json | 61 +++ IEclaude/data/__init__.py | 17 + IEclaude/data/traffic_dataloader.py | 415 ++++++++++++++ IEclaude/evaluate.py | 360 +++++++++++++ IEclaude/generate_configs.py | 83 +++ IEclaude/models/__init__.py | 18 + IEclaude/models/ie_diffusion.py | 652 ++++++++++++++++++++++ IEclaude/models/s4_explicit.py | 550 +++++++++++++++++++ IEclaude/models/tcn_implicit.py | 391 ++++++++++++++ IEclaude/requirements.txt | 27 + IEclaude/run_all.sh | 169 ++++++ IEclaude/train.py | 425 +++++++++++++++ IEclaude/utils/__init__.py | 12 + IEclaude/utils/diffusion.py | 365 +++++++++++++ 15 files changed, 4348 insertions(+) create mode 100644 IEclaude/README.md create mode 100644 IEclaude/configs/config_20.json create mode 100644 IEclaude/data/__init__.py create mode 100644 IEclaude/data/traffic_dataloader.py create mode 100644 IEclaude/evaluate.py create mode 100755 IEclaude/generate_configs.py create mode 100644 IEclaude/models/__init__.py create mode 100644 IEclaude/models/ie_diffusion.py create mode 100644 IEclaude/models/s4_explicit.py create mode 100644 IEclaude/models/tcn_implicit.py create mode 100644 IEclaude/requirements.txt create mode 100755 IEclaude/run_all.sh create mode 100644 IEclaude/train.py create mode 100644 IEclaude/utils/__init__.py create mode 100644 IEclaude/utils/diffusion.py diff --git a/IEclaude/README.md b/IEclaude/README.md new file mode 100644 index 0000000..2ab9e14 --- /dev/null +++ b/IEclaude/README.md @@ -0,0 +1,803 @@ +# IEclaude: 隐式显式扩散模型用于交通数据插补 + +**IEclaude** (Implicit-Explicit Claude) 是一个基于扩散模型的时间序列插补框架，结合了TCN隐式特征提取和S4显式特征提取，专门用于交通流量数据的缺失值插补任务。 + +--- + +## 目录 + +1. [项目简介](#项目简介) +2. [核心创新](#核心创新) +3. [架构设计](#架构设计) +4. [代码结构](#代码结构) +5. [详细代码解释](#详细代码解释) +6. [超参数调优指南](#超参数调优指南) +7. [快速开始](#快速开始) +8. [实验结果](#实验结果) +9. [常见问题](#常见问题) + +--- + +## 项目简介 + +### 背景 + +时间序列数据的缺失值问题在交通、气象、医疗等领域广泛存在。传统的插补方法（如线性插值、KNN）难以捕获复杂的时序依赖关系。本项目提出的**隐式显式扩散模型**通过以下方式解决这个问题： + +1. **扩散模型**: 基于DDPM的生成框架，通过逐步去噪实现高质量插补 +2. **隐式特征提取**: TCN捕获多时间尺度的局部依赖 +3. **显式特征提取**: S4捕获长期的全局依赖 + +### 论文参考 + +- **SSSD**: Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models +- **Traffic Diffusion**: A Diffusion Model for Traffic Data Imputation + +--- + +## 核心创新 + +### 1. 隐式特征提取（TCN） + +**TCN (Temporal Convolutional Network)** 使用扩张因果卷积提取不同时间尺度的隐式特征。 + +**核心优势**： +- 因果性：保证不使用未来信息 +- 多尺度：不同扩张率捕获不同时间尺度 +- 高效：并行计算，比RNN快 + +**实现细节**： +```python +# 扩张率：[1, 2, 4, 8] +# 感受野计算：RF = 1 + 2 * (kernel_size - 1) * sum(dilation_rates) +# 例如 kernel_size=3, dilation=[1,2,4,8] +# RF = 1 + 2*2*(1+2+4+8) = 61 时间步 +``` + +### 2. 显式特征提取（S4） + +**S4 (Structured State Space Model)** 使用状态空间模型捕获长期依赖。 + +**核心优势**： +- 长期记忆：HiPPO初始化保证长期信息保留 +- 线性复杂度：通过FFT加速到O(L log L) +- 理论保证：状态空间理论提供数学基础 + +**状态空间方程**： +``` +x_{t+1} = A * x_t + B * u_t (状态更新) +y_t = C * x_t + D * u_t (输出) +``` + +**ABCD矩阵说明**： +- **A矩阵**: 状态转移，控制历史如何演化（HiPPO初始化） +- **B矩阵**: 输入权重，控制当前输入的影响 +- **C矩阵**: 输出权重，控制状态到输出的映射 +- **D矩阵**: 跳跃连接，允许输入直接影响输出 + +### 3. 扩散模型 + +**DDPM (Denoising Diffusion Probabilistic Model)** 通过逐步去噪实现插补。 + +**前向过程（加噪）**： +``` +x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon +``` + +**反向过程（去噪）**： +``` +x_{t-1} = (1/sqrt(alpha_t)) * (x_t - (beta_t/sqrt(1-alpha_bar_t)) * epsilon_theta) +``` + +**关键参数**： +- **T**: 扩散步数（200），步数越多越平滑但越慢 +- **beta_0/beta_T**: 噪声调度（0.0001/0.02），控制噪声增长速度 + +--- + +## 架构设计 + +### 整体架构 + +``` +输入数据 (observed_data, mask) + ↓ +┌─────────────────────────────────────┐ +│ DETACH模块: 分解输入 │ +│ - 观测数据 (observed_data) │ +│ - 缺失mask (mask) │ +│ - 噪声 (noise) │ +└─────────────────────────────────────┘ + ↓ ↓ +┌──────────────┐ ┌──────────────┐ +│ TCN隐式提取 │ │ S4显式提取 │ +│ 多尺度特征 │ │ 长期依赖 │ +└──────────────┘ └──────────────┘ + ↓ ↓ + └────────┬───────────┘ + ↓ +┌─────────────────────────────────────┐ +│ 特征融合 + 条件信息 │ +│ [observed, mask, implicit, explicit]│ +└─────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────┐ +│ 残差块 x N │ +│ - 扩散步骤嵌入 │ +│ - 门控激活 │ +│ - 跳跃连接 │ +└─────────────────────────────────────┘ + ↓ + 去噪输出 +``` + +### 数据流 + +1. **输入**: `[B, C, L]` - 批次大小、通道数、序列长度 +2. **TCN输出**: `[B, 1, L]` - 隐式特征（单通道汇总） +3. **S4输出**: `[B, res_channels, L]` - 显式特征（多通道） +4. **条件信息**: `[B, C+C+1+res_channels, L]` - 拼接所有条件 +5. **最终输出**: `[B, C, L]` - 预测的噪声 + +--- + +## 代码结构 + +``` +IEclaude/ +├── models/ # 模型定义 +│ ├── __init__.py # 模块导出 +│ ├── tcn_implicit.py # TCN隐式特征提取 +│ ├── s4_explicit.py # S4显式特征提取 +│ └── ie_diffusion.py # 隐式显式扩散模型 +│ +├── data/ # 数据处理 +│ ├── __init__.py +│ └── traffic_dataloader.py # 数据加载和预处理 +│ +├── utils/ # 工具函数 +│ ├── __init__.py +│ └── diffusion.py # 扩散过程实现 +│ +├── configs/ # 配置文件 +│ ├── config_20.json # 20%缺失率配置 +│ ├── config_30.json # 30%缺失率配置 +│ └── ... # 其他缺失率 +│ +├── scripts/ # 实用脚本 +│ +├── results/ # 实验结果 +│ +├── train.py # 训练脚本 +├── evaluate.py # 评估脚本 +├── generate_configs.py # 生成配置文件 +├── run_all.sh # 一键运行所有实验 +├── requirements.txt # 依赖包 +└── README.md # 本文档 +``` + +--- + +## 详细代码解释 + +### 1. TCN隐式特征提取 (`models/tcn_implicit.py`) + +#### CausalConv1d 类 + +**作用**: 因果卷积，保证输出只依赖当前和过去的输入。 + +**核心代码**： +```python +# 计算padding: 左侧填充，右侧不填充 +self.padding = (kernel_size - 1) * dilation + +# 卷积前填充 +x = F.pad(x, (self.padding, 0)) # 只在左侧填充 + +# 执行卷积 +x = self.conv(x) +``` + +**调优参数**： +- `dilation`: 扩张率，控制感受野大小 + - dilation=1: 连续采样 + - dilation=2: 每隔1个位置采样 + - dilation=4: 每隔3个位置采样 + +#### TCNResidualBlock 类 + +**作用**: 残差块，允许梯度直接传播。 + +**核心代码**： +```python +# 两个卷积层 + 残差连接 +residual = x +out = self.relu(self.conv1(x)) +out = self.relu(self.conv2(out)) +out = out + residual # 残差连接 +``` + +**为什么使用残差**： +1. 解决梯度消失问题 +2. 允许网络学习增量（residual）而不是完整变换 +3. 使深层网络训练更稳定 + +#### TCNImplicitExtractor 类 + +**作用**: 完整的TCN特征提取器。 + +**关键参数**： +```python +hidden_channels=[256, 256, 256] # 各层通道数 +dilation_rates=[1, 2, 4, 8] # 扩张率序列 +``` + +**调优建议**： +1. **增加感受野**: 添加更大的扩张率，如`[1, 2, 4, 8, 16]` +2. **增加容量**: 增大通道数，如`[512, 512, 512]` +3. **防止过拟合**: 增大dropout，如`dropout=0.1` + +--- + +### 2. S4显式特征提取 (`models/s4_explicit.py`) + +#### HiPPO初始化 + +**作用**: 为A和B矩阵提供理论最优的初始化。 + +**核心代码**： +```python +def hippo_initializer(N): + # HiPPO-LegS矩阵 + A = np.zeros((N, N)) + for n in range(N): + for k in range(N): + if n > k: + A[n, k] = np.sqrt(2*n+1) * np.sqrt(2*k+1) + elif n == k: + A[n, k] = n + 1 + + B = np.sqrt(2 * np.arange(N) + 1).reshape(N, 1) + return A, B +``` + +**数学原理**： +- 基于Legendre多项式 +- 状态向量对应多项式系数 +- 能够高效压缩和记忆历史信息 + +#### 零阶保持器离散化 + +**作用**: 将连续系统离散化。 + +**公式**： +``` +A_bar = exp(A * dt) +B_bar = A^{-1} * (exp(A*dt) - I) * B +``` + +**参数dt**： +- 控制离散化步长 +- dt越小越精确但需要更多步数 +- 通常设为1.0 + +#### S4Layer 类 + +**核心流程**： +```python +1. 生成卷积核: K_l = C * A^l * B +2. FFT加速卷积: y = IFFT(FFT(x) * FFT(K)) +3. 添加skip connection: y = y + D * x +4. 激活和输出变换 +``` + +**调优参数**： +- `d_state (N)`: 状态维度 + - 增大：提升建模能力，但计算量增加 + - 建议：32-128 + +- `bidirectional`: 是否双向 + - True: 同时考虑过去和未来（插补任务推荐） + - False: 只考虑过去（预测任务推荐） + +--- + +### 3. 隐式显式扩散模型 (`models/ie_diffusion.py`) + +#### ResidualBlock 类 + +**作用**: 扩散模型的核心残差块。 + +**处理流程**： +```python +1. 添加扩散步骤嵌入 +h = x + fc_diffusion(diffusion_step_embed) + +2. 扩张卷积 +h = dilated_conv(h) + +3. 添加条件信息 +h = h + cond_conv(conditional) + +4. 门控激活 +h = tanh(h[:half]) * sigmoid(h[half:]) + +5. 残差和跳跃连接 +residual = res_conv(h) +skip = skip_conv(h) +return (x + residual) * sqrt(0.5), skip +``` + +**为什么使用门控激活**： +- 来自WaveNet +- tanh提供非线性，sigmoid提供门控 +- 能有效控制信息流 + +#### IEDiffusionModel 类 + +**初始化参数说明**： + +```python +# 基础参数 +in_channels: 输入通道数（传感器数量） +res_channels: 残差通道数（建议256-512） +num_res_layers: 残差层数（建议20-40） + +# TCN参数 +tcn_channels: TCN各层通道数 [256, 256, 256] +tcn_dilation_rates: 扩张率 [1, 2, 4, 8] + +# S4参数 +s4_d_state: S4状态维度（建议32-128） +s4_n_layers: S4层数（建议2-6） +s4_bidirectional: 是否双向（插补任务建议True） + +# 扩散参数 +diffusion_step_embed_dim_in/mid/out: 嵌入维度 +``` + +**前向传播流程**： +```python +1. TCN提取隐式特征: implicit = tcn_extractor(masked_data) +2. S4提取显式特征: explicit = s4_extractor(masked_data) +3. 拼接条件信息: cond = [masked_data, mask, implicit, explicit] +4. 初始化噪声: x = init_conv(noise) +5. 计算扩散嵌入: embed = fc(calc_diffusion_step_embedding(t)) +6. 通过残差块: + for block in residual_blocks: + x, skip = block(x, cond, embed) + skip_sum += skip +7. 输出: output = final_conv(skip_sum) +``` + +--- + +### 4. 扩散过程 (`utils/diffusion.py`) + +#### DiffusionProcess 类 + +**核心参数**： +```python +# 噪声调度参数 +betas: [beta_0, ..., beta_T] +alphas: 1 - betas +alpha_bars: cumprod(alphas) # 累积乘积 +``` + +**前向采样 q_sample**： +```python +# 从x_0直接采样x_t（不需要迭代） +x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise +``` + +**反向采样 p_sample**： +```python +# 从x_t采样x_{t-1}（需要模型预测） +predicted_noise = model(x_t, ...) +mean = (x_t - coef * predicted_noise) / sqrt(alpha_t) +x_{t-1} = mean + sqrt(posterior_var) * z # z~N(0,1) +``` + +**调优参数**： +- `T`: 扩散步数 + - 增大：质量提升但速度降低 + - 建议：200-500 + +- `beta_0/beta_T`: 噪声范围 + - beta_0: 起始噪声（建议0.0001-0.001） + - beta_T: 结束噪声（建议0.02-0.1） + +- `schedule`: 调度方式 + - 'linear': 线性增长（简单有效） + - 'cosine': 余弦增长（更平滑） + +--- + +### 5. 数据加载器 (`data/traffic_dataloader.py`) + +#### 数据预处理流程 + +```python +1. 读取数据: pd.read_csv(data_path) +2. 归一化: StandardScaler() or MinMaxScaler() +3. 切分序列: 滑动窗口 [seq_len] with stride +4. 划分训练/测试集: 按时间顺序 +5. 生成mask: 根据缺失率随机生成 +``` + +#### TrafficDataset 类 + +**功能**： +- 存储数据 +- 动态生成缺失mask +- 返回训练样本 + +**缺失模式**： +```python +'random': 每个位置独立地随机缺失 +'block': 连续的时间段缺失 +'spatial': 某些传感器整体缺失 +``` + +--- + +## 超参数调优指南 + +### 三个层次的调优 + +#### 1. 基础超参数（初学者） + +这些参数控制模型的基本容量和训练过程，不涉及复杂的架构调整。 + +**学习率 (learning_rate)** +```json +"learning_rate": 0.0002 +``` +- 作用：控制参数更新的步长 +- 调优建议： + - 训练不稳定 → 减小（0.0001） + - 训练太慢 → 增大（0.0005） + - 使用余弦调度自动调整 + +**批次大小 (batch_size)** +```json +"batch_size": 8 +``` +- 作用：每次训练使用的样本数 +- 调优建议： + - 显存不足 → 减小（4） + - 训练太慢 → 增大（16, 32） + - 建议：2的幂次 + +**训练轮数 (epochs)** +```json +"epochs": 100 +``` +- 作用：完整遍历训练集的次数 +- 调优建议： + - 观察训练曲线，loss不再下降时停止 + - 使用early stopping + +**Dropout** +```json +"tcn_dropout": 0.0, +"s4_dropout": 0.0 +``` +- 作用：防止过拟合 +- 调优建议： + - 过拟合（训练集好，测试集差） → 增大（0.1-0.2） + - 欠拟合 → 减小或设为0 + +#### 2. 架构调整（进阶） + +这些参数控制TCN和S4模块的结构，影响特征提取能力。 + +**TCN扩张率 (tcn_dilation_rates)** +```json +"tcn_dilation_rates": [1, 2, 4, 8] +``` +- 作用：控制TCN的感受野和时间尺度 +- 感受野计算：`RF = 1 + 2*(kernel_size-1)*sum(dilation_rates)` +- 调优建议： + - **标准配置**: `[1, 2, 4, 8]` - RF=61，适合大多数情况 + - **更长依赖**: `[1, 2, 4, 8, 16]` - RF=93，捕获更长依赖 + - **跳跃模式**: `[1, 4, 16]` - RF=63，稀疏采样 + +**TCN通道数 (tcn_channels)** +```json +"tcn_channels": [256, 256, 256] +``` +- 作用：控制TCN的表达能力 +- 调优建议： + - **标准**: `[256, 256, 256]` + - **更强**: `[512, 512, 512]` - 更强表达，更多参数 + - **递增**: `[128, 256, 512]` - 逐层提取更抽象特征 + +**S4状态维度 (s4_d_state)** +```json +"s4_d_state": 64 +``` +- 作用：S4的状态空间维度N，控制记忆容量 +- ABCD矩阵关系： + - A矩阵：[N, N] - 状态转移 + - B矩阵：[N, 1] - 输入映射 + - C矩阵：[1, N] - 输出映射 + - D矩阵：标量 - skip connection +- 调优建议： + - **小**: 32 - 快速但容量有限 + - **中**: 64 - 平衡点（推荐） + - **大**: 128 - 更强记忆但更慢 + +**S4层数 (s4_n_layers)** +```json +"s4_n_layers": 4 +``` +- 作用：S4的深度，更深提取更抽象特征 +- 调优建议： + - **浅**: 2 - 快速 + - **中**: 4 - 平衡（推荐） + - **深**: 6-8 - 更强但可能过拟合 + +**S4双向模式 (s4_bidirectional)** +```json +"s4_bidirectional": true +``` +- 作用：是否同时考虑过去和未来 +- 调优建议： + - **插补任务**: true - 利用全部信息 + - **预测任务**: false - 只用历史信息 + +#### 3. 扩散模型参数（高级） + +这些参数控制扩散过程，影响生成质量和速度。 + +**扩散步数 (T)** +```json +"T": 200 +``` +- 作用：扩散和去噪的总步数 +- 影响： + - 步数越多，过程越平滑，质量越好，但越慢 + - 步数越少，速度快，但质量可能下降 +- 调优建议： + - **快速测试**: 50-100 + - **标准**: 200 - 平衡点 + - **高质量**: 500-1000 + +**噪声调度 (beta_0, beta_T)** +```json +"beta_0": 0.0001, +"beta_T": 0.02 +``` +- 作用：控制噪声增长速度 +- 数学含义： + - beta_t：第t步的噪声水平 + - alpha_t = 1 - beta_t + - alpha_bar_t = prod(alpha_1, ..., alpha_t) +- 调优建议： + - **训练不稳定**: 减小beta_T（0.01） + - **生成质量差**: 调整beta范围 + - **标准配置**: beta_0=0.0001, beta_T=0.02 + +**噪声调度方式 (schedule)** +```json +"schedule": "linear" +``` +- 'linear': beta线性增长 +- 'cosine': beta按余弦曲线增长（更平滑） + +**扩散步骤嵌入维度 (embed_dim_in/mid/out)** +```json +"embed_dim_in": 128, +"embed_dim_mid": 512, +"embed_dim_out": 512 +``` +- 作用：将扩散步骤t编码为高维向量 +- 调优建议：通常不需要调整 + +--- + +### 调优流程建议 + +#### 第一阶段：基础调优 +1. 使用默认架构参数 +2. 只调整learning_rate和batch_size +3. 观察训练曲线是否平稳 +4. 目标：稳定训练 + +#### 第二阶段：架构调优 +1. 实验不同的TCN扩张率 + - 对比 [1,2,4,8] vs [1,2,4,8,16] vs [1,4,16] +2. 调整S4状态维度 + - 对比 32 vs 64 vs 128 +3. 调整模型深度 + - num_res_layers: 20 vs 36 vs 50 + - s4_n_layers: 2 vs 4 vs 6 +4. 目标：找到最佳架构 + +#### 第三阶段：精细调优 +1. 调整扩散步数T +2. 尝试不同的噪声调度 +3. 微调学习率和正则化 +4. 目标：达到目标指标 + +--- + +## 快速开始 + +### 环境准备 + +```bash +# 1. 克隆项目 +cd /home/user/SSSD/IEclaude + +# 2. 创建conda环境 +conda create -n ieclaude python=3.8 +conda activate ieclaude + +# 3. 安装依赖 +pip install -r requirements.txt + +# 4. 安装PyTorch（根据您的CUDA版本） +# CUDA 11.3 +pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 +``` + +### 数据准备 + +确保数据集位于正确位置： +```bash +# 检查数据文件 +ls /home/zhu/sssdtcn/LD2011_2014.txt +``` + +### 快速训练单个模型 + +```bash +# 1. 生成配置文件 +python generate_configs.py + +# 2. 训练20%缺失率的模型 +python train.py --config configs/config_20.json --gpu 0 + +# 3. 评估模型 +python evaluate.py --config configs/config_20.json --checkpoint results/traffic_20/best_model.pt --gpu 0 +``` + +### 运行所有实验 + +```bash +# 一键运行所有缺失率（20%-80%）的训练和评估 +chmod +x run_all.sh +./run_all.sh --gpu 0 +``` + +--- + +## 实验结果 + +### 目标指标 + +| 缺失率 | MAE | RMSE | +|--------|--------|--------| +| 20% | 0.272 | 0.389 | +| 30% | 0.297 | 0.424 | +| 40% | 0.334 | 0.477 | +| 50% | 0.378 | 0.540 | +| 60% | 0.450 | 0.655 | +| 70% | 0.541 | 0.776 | +| 80% | 0.732 | 1.049 | + +### 结果分析 + +所有实验结果保存在 `results/` 目录： + +``` +results/ +├── traffic_20/ +│ ├── best_model.pt # 最佳模型 +│ ├── training_loss.png # 训练曲线 +│ ├── evaluation_metrics.txt # 评估指标 +│ └── imputation_visualization.png # 插补可视化 +├── traffic_30/ +│ └── ... +└── final_report.txt # 汇总报告 +``` + +--- + +## 常见问题 + +### Q1: 显存不足怎么办？ + +**解决方案**： +1. 减小batch_size（8 → 4 → 2） +2. 减小模型尺寸： + ```json + "res_channels": 128, // 256 → 128 + "num_res_layers": 20 // 36 → 20 + ``` +3. 使用梯度累积 + +### Q2: 训练速度太慢？ + +**解决方案**： +1. 减小扩散步数T（200 → 100） +2. 使用单向S4（bidirectional: false） +3. 减少评估频率 +4. 使用更强的GPU + +### Q3: 指标不如预期？ + +**调优建议**： +1. 增加训练轮数（100 → 200） +2. 调整学习率（0.0002 → 0.0001） +3. 增大模型容量： + - res_channels: 256 → 512 + - num_res_layers: 36 → 50 +4. 尝试不同的TCN扩张率 + +### Q4: 训练不稳定，loss波动大？ + +**解决方案**： +1. 减小学习率（0.0002 → 0.0001） +2. 使用warm-up +3. 减小beta_T（0.02 → 0.01） +4. 增加batch_size + +### Q5: 如何修改序列长度？ + +**步骤**： +1. 修改配置文件： + ```json + "seq_len": 336 // 168 → 336 (2周) + ``` +2. 相应调整TCN感受野 +3. 更新s4_l_max + +### Q6: 可以用于其他数据集吗？ + +**是的！** 只需： +1. 修改 `data_path` +2. 调整 `seq_len` 和 `in_channels` +3. 相应调整模型参数 + +--- + +## 项目维护 + +### 代码规范 + +- 所有代码都有详细的中文注释 +- 每个函数都有docstring说明 +- 超参数都有调优指南 + +### 文件说明 + +- `train.py`: 训练主脚本 +- `evaluate.py`: 评估主脚本 +- `generate_configs.py`: 生成配置文件 +- `run_all.sh`: 一键运行所有实验 + +--- + +## 致谢 + +本项目参考了以下工作： + +1. **SSSD**: Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models +2. **S4**: Efficiently Modeling Long Sequences with Structured State Spaces +3. **DDPM**: Denoising Diffusion Probabilistic Models +4. **TCN**: An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling + +--- + +## 联系方式 + +如有问题或建议，请通过以下方式联系： +- GitHub Issues +- Email: [your-email] + +--- + +**祝您使用愉快！Good Luck! 🚀** diff --git a/IEclaude/configs/config_20.json b/IEclaude/configs/config_20.json new file mode 100644 index 0000000..6b9329c --- /dev/null +++ b/IEclaude/configs/config_20.json @@ -0,0 +1,61 @@ +{ + "comment": "IEclaude配置文件 - 20%缺失率", + + "seed": 42, + + "data": { + "data_path": "/home/zhu/sssdtcn/LD2011_2014.txt", + "seq_len": 168, + "stride": 84, + "train_ratio": 0.7, + "normalize": "standard" + }, + + "model": { + "res_channels": 256, + "skip_channels": 256, + "num_res_layers": 36, + "dilation_cycle": 10, + + "tcn_channels": [256, 256, 256], + "tcn_kernel_size": 3, + "tcn_dilation_rates": [1, 2, 4, 8], + "tcn_dropout": 0.0, + + "s4_d_state": 64, + "s4_n_layers": 4, + "s4_dropout": 0.0, + "s4_bidirectional": true + }, + + "diffusion": { + "T": 200, + "beta_0": 0.0001, + "beta_T": 0.02, + "schedule": "linear", + + "embed_dim_in": 128, + "embed_dim_mid": 512, + "embed_dim_out": 512 + }, + + "train": { + "output_dir": "./results/traffic_20", + "missing_rate": 0.2, + "missing_pattern": "random", + + "batch_size": 8, + "epochs": 100, + "learning_rate": 0.0002, + "weight_decay": 0.0, + + "scheduler": "cosine", + "min_lr": 1e-6, + + "save_interval": 10, + "num_workers": 0, + + "only_generate_missing": true, + "clean_before_train": true + } +} diff --git a/IEclaude/data/__init__.py b/IEclaude/data/__init__.py new file mode 100644 index 0000000..913de14 --- /dev/null +++ b/IEclaude/data/__init__.py @@ -0,0 +1,17 @@ +""" +IEclaude Data Module + +包含数据加载和预处理工具 +""" + +from .traffic_dataloader import ( + TrafficDataset, + load_traffic_data, + create_dataloader +) + +__all__ = [ + 'TrafficDataset', + 'load_traffic_data', + 'create_dataloader' +] diff --git a/IEclaude/data/traffic_dataloader.py b/IEclaude/data/traffic_dataloader.py new file mode 100644 index 0000000..975ec0a --- /dev/null +++ b/IEclaude/data/traffic_dataloader.py @@ -0,0 +1,415 @@ +""" +交通数据加载器 (Traffic Data Loader) + +功能说明： + 加载和预处理LD2011_2014.txt交通数据集，用于时间序列插补任务 + +数据集说明： + - LD2011_2014.txt: 电力负载数据集 + - 包含370个客户的电力消耗时间序列 + - 采样频率: 每15分钟一个数据点 + - 时间跨度: 2011-2014年 + +预处理流程： + 1. 读取原始数据 + 2. 数据归一化（标准化或最小-最大缩放） + 3. 切分为固定长度的序列段 + 4. 生成训练/测试集 + 5. 创建缺失mask（根据指定缺失率） + +超参数调节指南： + 1. 数据处理参数： + - seq_len: 序列长度 [建议: 168 for 1周] + - stride: 滑动窗口步长 [建议: seq_len//2] + - train_ratio: 训练集比例 [建议: 0.7-0.8] + - normalize_method: 归一化方法 ['standard', 'minmax'] + + 2. 缺失模式： + - missing_rate: 缺失率 [0.2-0.8] + - missing_pattern: 缺失模式 ['random', 'block', 'spatial'] + * random: 随机缺失 + * block: 连续块缺失 + * spatial: 空间缺失（某些传感器整体缺失） + + 3. 数据增强： + - add_noise: 是否添加噪声 + - noise_std: 噪声标准差 +""" + +import numpy as np +import torch +from torch.utils.data import Dataset, DataLoader +import pandas as pd +from sklearn.preprocessing import StandardScaler, MinMaxScaler +import os + + +class TrafficDataset(Dataset): + """ + 交通数据集类 + + 功能： + - 存储预处理后的数据 + - 生成缺失mask + - 返回训练样本 + + 参数： + data (np.array): 数据 [num_samples, num_channels, seq_len] + missing_rate (float): 缺失率 + missing_pattern (str): 缺失模式 + """ + def __init__(self, data, missing_rate=0.3, missing_pattern='random'): + super().__init__() + + self.data = torch.FloatTensor(data) # [N, C, L] + self.missing_rate = missing_rate + self.missing_pattern = missing_pattern + + self.num_samples, self.num_channels, self.seq_len = self.data.shape + + def __len__(self): + return self.num_samples + + def __getitem__(self, idx): + """ + 获取一个样本 + + 返回： + sample (dict): 包含以下键值 + - 'observed_data': 观测数据 [C, L] + - 'mask': 缺失mask [C, L], 1=观测，0=缺失 + - 'ground_truth': 真实数据 [C, L] + - 'index': 样本索引 + """ + # 获取原始数据 + ground_truth = self.data[idx] # [C, L] + + # 生成缺失mask + mask = self._generate_mask(ground_truth.shape) + + # 应用mask生成观测数据 + observed_data = ground_truth * mask + + sample = { + 'observed_data': observed_data, + 'mask': mask, + 'ground_truth': ground_truth, + 'index': idx + } + + return sample + + def _generate_mask(self, shape): + """ + 生成缺失mask + + 参数： + shape (tuple): 数据形状 [C, L] + + 返回： + mask (torch.Tensor): 缺失mask [C, L] + """ + C, L = shape + + if self.missing_pattern == 'random': + # 随机缺失：每个位置独立地以missing_rate的概率缺失 + mask = torch.rand(C, L) > self.missing_rate + + elif self.missing_pattern == 'block': + # 块缺失：连续的时间段缺失 + mask = torch.ones(C, L) + # 随机选择几个连续块进行缺失 + num_blocks = int(self.missing_rate * L / 10) # 每个块大约10个时间步 + for _ in range(num_blocks): + for c in range(C): + block_start = np.random.randint(0, L - 10) + block_len = np.random.randint(5, 15) + block_end = min(block_start + block_len, L) + mask[c, block_start:block_end] = 0 + + elif self.missing_pattern == 'spatial': + # 空间缺失：某些通道完全缺失 + mask = torch.ones(C, L) + num_missing_channels = int(self.missing_rate * C) + missing_channels = np.random.choice(C, num_missing_channels, replace=False) + mask[missing_channels, :] = 0 + + else: + raise ValueError(f"未知的缺失模式: {self.missing_pattern}") + + return mask.float() + + +def load_traffic_data( + data_path, + seq_len=168, + stride=84, + train_ratio=0.7, + normalize_method='standard', + verbose=True +): + """ + 加载交通数据 + + 参数： + data_path (str): 数据文件路径 + seq_len (int): 序列长度（时间步数） + stride (int): 滑动窗口步长 + train_ratio (float): 训练集比例 + normalize_method (str): 归一化方法 ['standard', 'minmax', 'none'] + verbose (bool): 是否打印信息 + + 返回： + train_data (np.array): 训练数据 [N_train, C, L] + test_data (np.array): 测试数据 [N_test, C, L] + scaler: 归一化器（用于反归一化） + data_info (dict): 数据信息 + """ + if verbose: + print("=" * 80) + print("加载交通数据集") + print("=" * 80) + + # ======================================================================== + # 步骤1: 读取原始数据 + # ======================================================================== + if verbose: + print(f"\n1. 读取数据文件: {data_path}") + + if not os.path.exists(data_path): + raise FileNotFoundError(f"数据文件不存在: {data_path}") + + # 读取数据 + # LD2011_2014.txt是以分号分隔的CSV文件 + # 第一列是时间戳，其余370列是客户数据 + try: + # 尝试读取带分号的CSV + df = pd.read_csv(data_path, sep=';', decimal=',') + except: + # 如果失败，尝试普通CSV + df = pd.read_csv(data_path) + + if verbose: + print(f" 原始数据形状: {df.shape}") + print(f" 列名: {df.columns.tolist()[:5]}...") # 显示前5个列名 + + # 移除第一列（时间戳或索引） + # 假设第一列不是数据列 + if df.shape[1] == 371: # 1个时间戳 + 370个客户 + data = df.iloc[:, 1:].values.astype(np.float32) + else: + data = df.values.astype(np.float32) + + if verbose: + print(f" 数据数组形状: {data.shape} [时间步, 通道数]") + + # ======================================================================== + # 步骤2: 数据归一化 + # ======================================================================== + if verbose: + print(f"\n2. 数据归一化 (方法: {normalize_method})") + + if normalize_method == 'standard': + # 标准化：减去均值，除以标准差 + # 对每个通道独立标准化 + scaler = StandardScaler() + data = scaler.fit_transform(data) # [T, C] + + elif normalize_method == 'minmax': + # 最小-最大缩放：缩放到[0, 1] + scaler = MinMaxScaler() + data = scaler.fit_transform(data) # [T, C] + + elif normalize_method == 'none': + # 不进行归一化 + scaler = None + + else: + raise ValueError(f"未知的归一化方法: {normalize_method}") + + if verbose: + print(f" 归一化后统计:") + print(f" 均值: {data.mean():.4f}") + print(f" 标准差: {data.std():.4f}") + print(f" 最小值: {data.min():.4f}") + print(f" 最大值: {data.max():.4f}") + + # ======================================================================== + # 步骤3: 切分为固定长度的序列段 + # ======================================================================== + if verbose: + print(f"\n3. 切分序列 (seq_len={seq_len}, stride={stride})") + + T, C = data.shape # T: 时间步数, C: 通道数（传感器数量） + + # 使用滑动窗口切分 + sequences = [] + for start in range(0, T - seq_len + 1, stride): + end = start + seq_len + seq = data[start:end, :] # [seq_len, C] + # 转置为 [C, seq_len] 格式（PyTorch卷积的标准格式） + seq = seq.T # [C, seq_len] + sequences.append(seq) + + sequences = np.array(sequences) # [N, C, seq_len] + + if verbose: + print(f" 序列数量: {len(sequences)}") + print(f" 序列形状: {sequences.shape} [样本数, 通道数, 序列长度]") + + # ======================================================================== + # 步骤4: 切分训练集和测试集 + # ======================================================================== + if verbose: + print(f"\n4. 切分训练/测试集 (训练比例: {train_ratio})") + + num_samples = len(sequences) + num_train = int(num_samples * train_ratio) + + # 按时间顺序划分：前面的用于训练，后面的用于测试 + train_data = sequences[:num_train] + test_data = sequences[num_train:] + + if verbose: + print(f" 训练集: {train_data.shape}") + print(f" 测试集: {test_data.shape}") + + # ======================================================================== + # 步骤5: 数据信息 + # ======================================================================== + data_info = { + 'num_channels': C, + 'seq_len': seq_len, + 'num_train': len(train_data), + 'num_test': len(test_data), + 'normalize_method': normalize_method, + 'original_shape': df.shape, + 'train_ratio': train_ratio, + 'stride': stride + } + + if verbose: + print("\n" + "=" * 80) + print("数据加载完成！") + print("=" * 80) + + return train_data, test_data, scaler, data_info + + +def create_dataloader( + data, + missing_rate=0.3, + missing_pattern='random', + batch_size=8, + shuffle=True, + num_workers=0 +): + """ + 创建数据加载器 + + 参数： + data (np.array): 数据 [N, C, L] + missing_rate (float): 缺失率 + missing_pattern (str): 缺失模式 + batch_size (int): 批次大小 + shuffle (bool): 是否打乱数据 + num_workers (int): 数据加载线程数 + + 返回： + dataloader (DataLoader): PyTorch数据加载器 + """ + dataset = TrafficDataset( + data=data, + missing_rate=missing_rate, + missing_pattern=missing_pattern + ) + + dataloader = DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=True # 加速GPU传输 + ) + + return dataloader + + +# ============================================================================ +# 使用示例和调试代码 +# ============================================================================ + +if __name__ == "__main__": + """ + 测试数据加载器 + """ + print("=" * 80) + print("交通数据加载器测试") + print("=" * 80) + + # 数据文件路径（实际使用时需要修改为真实路径） + data_path = "/home/zhu/sssdtcn/LD2011_2014.txt" + + # 检查文件是否存在 + if not os.path.exists(data_path): + print(f"\n警告: 数据文件不存在: {data_path}") + print("测试将使用模拟数据") + + # 创建模拟数据用于测试 + print("\n创建模拟数据...") + num_timesteps = 10000 + num_channels = 370 + mock_data = np.random.randn(num_timesteps, num_channels).astype(np.float32) + + # 保存为临时文件 + import tempfile + with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f: + data_path = f.name + # 写入CSV格式 + pd.DataFrame(mock_data).to_csv(f, index=False) + + print(f"模拟数据已保存到: {data_path}") + + # 加载数据 + train_data, test_data, scaler, data_info = load_traffic_data( + data_path=data_path, + seq_len=168, + stride=84, + train_ratio=0.7, + normalize_method='standard', + verbose=True + ) + + print(f"\n数据信息:") + for key, value in data_info.items(): + print(f" {key}: {value}") + + # 创建数据加载器 + print(f"\n创建数据加载器...") + train_loader = create_dataloader( + data=train_data, + missing_rate=0.3, + missing_pattern='random', + batch_size=8, + shuffle=True + ) + + print(f" 训练集batch数量: {len(train_loader)}") + + # 测试获取一个batch + print(f"\n获取一个batch测试:") + for batch in train_loader: + print(f" observed_data: {batch['observed_data'].shape}") + print(f" mask: {batch['mask'].shape}") + print(f" ground_truth: {batch['ground_truth'].shape}") + print(f" 实际缺失率: {1 - batch['mask'].mean():.2%}") + break + + print("\n" + "=" * 80) + print("测试完成！") + print("=" * 80) + + # 清理临时文件 + if 'mock_data' in locals(): + os.unlink(data_path) + print(f"\n已清理临时文件: {data_path}") diff --git a/IEclaude/evaluate.py b/IEclaude/evaluate.py new file mode 100644 index 0000000..d3ec562 --- /dev/null +++ b/IEclaude/evaluate.py @@ -0,0 +1,360 @@ +""" +IEclaude评估脚本 (Evaluation Script) + +功能： + 加载训练好的模型，进行插补并评估性能 + +使用方法： + python evaluate.py --config configs/config_20.json --checkpoint results/traffic_20/best_model.pt --gpu 0 + +评估指标： + - MAE (Mean Absolute Error): 平均绝对误差 + - RMSE (Root Mean Squared Error): 均方根误差 +""" + +import os +import sys +import json +import argparse +import torch +import numpy as np +import matplotlib.pyplot as plt +from tqdm import tqdm + +# 导入模块 +from models import IEDiffusionModel +from data import load_traffic_data, create_dataloader +from utils import DiffusionProcess + + +def parse_args(): + """解析命令行参数""" + parser = argparse.ArgumentParser(description='评估IEclaude模型') + + parser.add_argument('--config', type=str, required=True, + help='配置文件路径') + parser.add_argument('--checkpoint', type=str, required=True, + help='模型checkpoint路径') + parser.add_argument('--gpu', type=int, default=0, + help='GPU编号') + parser.add_argument('--num_samples', type=int, default=10, + help='每个样本的生成次数（用于平均）') + + return parser.parse_args() + + +def load_config(config_path): + """加载配置文件""" + with open(config_path, 'r', encoding='utf-8') as f: + config = json.load(f) + return config + + +def calculate_metrics(predictions, ground_truth, mask): + """ + 计算评估指标 + + 参数： + predictions (np.array): 预测值 [N, C, L] + ground_truth (np.array): 真实值 [N, C, L] + mask (np.array): mask [N, C, L], 1=观测, 0=缺失 + + 返回： + metrics (dict): 包含MAE和RMSE的字典 + """ + # 只计算缺失位置的指标 + missing_mask = (1 - mask).astype(bool) + + # 提取缺失位置的预测值和真实值 + pred_missing = predictions[missing_mask] + true_missing = ground_truth[missing_mask] + + # 计算MAE + mae = np.mean(np.abs(pred_missing - true_missing)) + + # 计算RMSE + mse = np.mean((pred_missing - true_missing) ** 2) + rmse = np.sqrt(mse) + + metrics = { + 'MAE': mae, + 'RMSE': rmse, + 'num_missing': len(pred_missing) + } + + return metrics + + +def impute(model, dataloader, diffusion_process, device, num_samples=10, verbose=True): + """ + 执行插补 + + 参数： + model: 模型 + dataloader: 数据加载器 + diffusion_process: 扩散过程 + device: 设备 + num_samples: 每个样本的生成次数 + verbose: 是否显示进度 + + 返回： + all_predictions (np.array): 所有预测值 + all_ground_truth (np.array): 所有真实值 + all_masks (np.array): 所有mask + """ + model.eval() + + all_predictions = [] + all_ground_truth = [] + all_masks = [] + + with torch.no_grad(): + pbar = tqdm(dataloader, desc='Imputing') if verbose else dataloader + + for batch in pbar: + observed_data = batch['observed_data'].to(device) + mask = batch['mask'].to(device) + ground_truth = batch['ground_truth'].to(device) + + B, C, L = observed_data.shape + + # 多次采样取平均（降低随机性） + samples = [] + for _ in range(num_samples): + # 使用扩散过程生成 + imputed = diffusion_process.p_sample_loop( + model=model, + shape=(B, C, L), + observed_data=observed_data, + mask=mask, + device=device, + verbose=False + ) + samples.append(imputed.cpu().numpy()) + + # 平均多次采样的结果 + imputed_avg = np.mean(samples, axis=0) + + # 合并观测值和插补值 + final = observed_data.cpu().numpy() * mask.cpu().numpy() + \ + imputed_avg * (1 - mask.cpu().numpy()) + + all_predictions.append(final) + all_ground_truth.append(ground_truth.cpu().numpy()) + all_masks.append(mask.cpu().numpy()) + + # 拼接所有batch + all_predictions = np.concatenate(all_predictions, axis=0) + all_ground_truth = np.concatenate(all_ground_truth, axis=0) + all_masks = np.concatenate(all_masks, axis=0) + + return all_predictions, all_ground_truth, all_masks + + +def visualize_imputation(predictions, ground_truth, mask, save_path, num_examples=5): + """ + 可视化插补结果 + + 参数： + predictions: 预测值 [N, C, L] + ground_truth: 真实值 [N, C, L] + mask: mask [N, C, L] + save_path: 保存路径 + num_examples: 显示的样本数量 + """ + num_examples = min(num_examples, len(predictions)) + + fig, axes = plt.subplots(num_examples, 1, figsize=(15, 3 * num_examples)) + if num_examples == 1: + axes = [axes] + + for i in range(num_examples): + # 选择一个通道进行可视化（第一个通道） + pred = predictions[i, 0, :] + true = ground_truth[i, 0, :] + m = mask[i, 0, :] + + # 创建x轴 + x = np.arange(len(pred)) + + # 绘制真实值 + axes[i].plot(x, true, 'k-', label='Ground Truth', linewidth=1.5, alpha=0.7) + + # 绘制观测值 + observed_indices = m == 1 + axes[i].scatter(x[observed_indices], true[observed_indices], + c='blue', s=20, label='Observed', zorder=3) + + # 绘制插补值 + missing_indices = m == 0 + axes[i].scatter(x[missing_indices], pred[missing_indices], + c='red', s=20, label='Imputed', zorder=3) + + axes[i].set_xlabel('Time Step') + axes[i].set_ylabel('Value') + axes[i].set_title(f'Sample {i+1} - Channel 0') + axes[i].legend() + axes[i].grid(True, alpha=0.3) + + plt.tight_layout() + plt.savefig(save_path, dpi=300, bbox_inches='tight') + plt.close() + + +def evaluate(config, args): + """主评估函数""" + print("=" * 80) + print("IEclaude 评估脚本") + print("=" * 80) + + # 设置设备 + device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') + print(f"\n使用设备: {device}") + + # 输出目录 + output_dir = config['train']['output_dir'] + print(f"输出目录: {output_dir}") + + # ======================================================================== + # 1. 加载数据 + # ======================================================================== + print("\n" + "=" * 80) + print("1. 加载数据") + print("=" * 80) + + train_data, test_data, scaler, data_info = load_traffic_data( + data_path=config['data']['data_path'], + seq_len=config['data']['seq_len'], + stride=config['data'].get('stride', config['data']['seq_len'] // 2), + train_ratio=config['data'].get('train_ratio', 0.7), + normalize_method=config['data'].get('normalize', 'standard'), + verbose=True + ) + + # 创建测试数据加载器 + test_loader = create_dataloader( + data=test_data, + missing_rate=config['train']['missing_rate'], + missing_pattern=config['train'].get('missing_pattern', 'random'), + batch_size=config['train']['batch_size'], + shuffle=False, + num_workers=0 + ) + + print(f"\n测试集batch数量: {len(test_loader)}") + + # ======================================================================== + # 2. 加载模型 + # ======================================================================== + print("\n" + "=" * 80) + print("2. 加载模型") + print("=" * 80) + + model = IEDiffusionModel( + in_channels=data_info['num_channels'], + res_channels=config['model']['res_channels'], + skip_channels=config['model']['skip_channels'], + out_channels=data_info['num_channels'], + num_res_layers=config['model']['num_res_layers'], + dilation_cycle=config['model'].get('dilation_cycle', 10), + diffusion_step_embed_dim_in=config['diffusion']['embed_dim_in'], + diffusion_step_embed_dim_mid=config['diffusion']['embed_dim_mid'], + diffusion_step_embed_dim_out=config['diffusion']['embed_dim_out'], + tcn_channels=config['model']['tcn_channels'], + tcn_kernel_size=config['model']['tcn_kernel_size'], + tcn_dilation_rates=config['model']['tcn_dilation_rates'], + tcn_dropout=config['model'].get('tcn_dropout', 0.0), + s4_d_state=config['model']['s4_d_state'], + s4_n_layers=config['model']['s4_n_layers'], + s4_l_max=config['data']['seq_len'], + s4_dropout=config['model'].get('s4_dropout', 0.0), + s4_bidirectional=config['model'].get('s4_bidirectional', True) + ).to(device) + + # 加载checkpoint + checkpoint = torch.load(args.checkpoint, map_location=device) + model.load_state_dict(checkpoint['model_state_dict']) + print(f"已加载checkpoint: {args.checkpoint}") + print(f" Epoch: {checkpoint.get('epoch', 'N/A')}") + print(f" Loss: {checkpoint.get('loss', 'N/A'):.6f}") + + # ======================================================================== + # 3. 创建扩散过程 + # ======================================================================== + print("\n" + "=" * 80) + print("3. 创建扩散过程") + print("=" * 80) + + diffusion_process = DiffusionProcess( + T=config['diffusion']['T'], + beta_0=config['diffusion']['beta_0'], + beta_T=config['diffusion']['beta_T'], + schedule=config['diffusion'].get('schedule', 'linear') + ) + + # ======================================================================== + # 4. 执行插补 + # ======================================================================== + print("\n" + "=" * 80) + print("4. 执行插补") + print("=" * 80) + + predictions, ground_truth, masks = impute( + model=model, + dataloader=test_loader, + diffusion_process=diffusion_process, + device=device, + num_samples=args.num_samples, + verbose=True + ) + + print(f"\n插补完成！") + print(f" 预测形状: {predictions.shape}") + + # ======================================================================== + # 5. 计算指标 + # ======================================================================== + print("\n" + "=" * 80) + print("5. 计算指标") + print("=" * 80) + + metrics = calculate_metrics(predictions, ground_truth, masks) + + print(f"\n评估指标 (缺失率={config['train']['missing_rate']*100:.0f}%):") + print(f" MAE: {metrics['MAE']:.3f}") + print(f" RMSE: {metrics['RMSE']:.3f}") + print(f" 缺失样本数: {metrics['num_missing']:,}") + + # 保存指标 + metrics_file = os.path.join(output_dir, 'evaluation_metrics.txt') + with open(metrics_file, 'w') as f: + f.write(f"Missing Rate: {config['train']['missing_rate']*100:.0f}%\n") + f.write(f"MAE: {metrics['MAE']:.3f}\n") + f.write(f"RMSE: {metrics['RMSE']:.3f}\n") + f.write(f"Num Missing: {metrics['num_missing']}\n") + + print(f"\n指标已保存: {metrics_file}") + + # ======================================================================== + # 6. 可视化 + # ======================================================================== + print("\n" + "=" * 80) + print("6. 生成可视化") + print("=" * 80) + + vis_path = os.path.join(output_dir, 'imputation_visualization.png') + visualize_imputation(predictions, ground_truth, masks, vis_path, num_examples=5) + print(f"可视化已保存: {vis_path}") + + # ======================================================================== + # 7. 完成 + # ======================================================================== + print("\n" + "=" * 80) + print("评估完成！") + print("=" * 80) + + +if __name__ == '__main__': + args = parse_args() + config = load_config(args.config) + evaluate(config, args) diff --git a/IEclaude/generate_configs.py b/IEclaude/generate_configs.py new file mode 100755 index 0000000..3f0c10b --- /dev/null +++ b/IEclaude/generate_configs.py @@ -0,0 +1,83 @@ +""" +生成所有缺失率的配置文件 + +使用方法： + python generate_configs.py +""" + +import json +import os + +# 基础配置 +base_config = { + "seed": 42, + "data": { + "data_path": "/home/zhu/sssdtcn/LD2011_2014.txt", + "seq_len": 168, + "stride": 84, + "train_ratio": 0.7, + "normalize": "standard" + }, + "model": { + "res_channels": 256, + "skip_channels": 256, + "num_res_layers": 36, + "dilation_cycle": 10, + "tcn_channels": [256, 256, 256], + "tcn_kernel_size": 3, + "tcn_dilation_rates": [1, 2, 4, 8], + "tcn_dropout": 0.0, + "s4_d_state": 64, + "s4_n_layers": 4, + "s4_dropout": 0.0, + "s4_bidirectional": True + }, + "diffusion": { + "T": 200, + "beta_0": 0.0001, + "beta_T": 0.02, + "schedule": "linear", + "embed_dim_in": 128, + "embed_dim_mid": 512, + "embed_dim_out": 512 + }, + "train": { + "batch_size": 8, + "epochs": 100, + "learning_rate": 0.0002, + "weight_decay": 0.0, + "scheduler": "cosine", + "min_lr": 1e-6, + "save_interval": 10, + "num_workers": 0, + "only_generate_missing": True, + "clean_before_train": True, + "missing_pattern": "random" + } +} + +# 创建configs目录 +os.makedirs('configs', exist_ok=True) + +# 生成不同缺失率的配置文件 +missing_rates = [20, 30, 40, 50, 60, 70, 80] + +for rate in missing_rates: + config = base_config.copy() + config = json.loads(json.dumps(config)) # 深拷贝 + + # 设置缺失率 + config['train']['missing_rate'] = rate / 100.0 + config['train']['output_dir'] = f'./results/traffic_{rate}' + + # 添加注释 + config['comment'] = f'IEclaude配置文件 - {rate}%缺失率' + + # 保存配置文件 + config_path = f'configs/config_{rate}.json' + with open(config_path, 'w', encoding='utf-8') as f: + json.dump(config, f, indent=4, ensure_ascii=False) + + print(f'已生成配置文件: {config_path}') + +print(f'\n共生成 {len(missing_rates)} 个配置文件') diff --git a/IEclaude/models/__init__.py b/IEclaude/models/__init__.py new file mode 100644 index 0000000..ae6b2a3 --- /dev/null +++ b/IEclaude/models/__init__.py @@ -0,0 +1,18 @@ +""" +IEclaude Models Module + +包含所有模型组件： +- TCN隐式特征提取 +- S4显式特征提取 +- 隐式显式扩散模型 +""" + +from .tcn_implicit import TCNImplicitExtractor +from .s4_explicit import S4ExplicitExtractor +from .ie_diffusion import IEDiffusionModel + +__all__ = [ + 'TCNImplicitExtractor', + 'S4ExplicitExtractor', + 'IEDiffusionModel' +] diff --git a/IEclaude/models/ie_diffusion.py b/IEclaude/models/ie_diffusion.py new file mode 100644 index 0000000..533159e --- /dev/null +++ b/IEclaude/models/ie_diffusion.py @@ -0,0 +1,652 @@ +""" +隐式显式扩散模型 (Implicit-Explicit Diffusion Model for Traffic Data Imputation) + +论文参考： + A Diffusion Model for Traffic Data Imputation + +功能说明： + 结合TCN隐式特征提取和S4显式特征提取的扩散模型，用于时间序列插补 + +核心创新： + 1. 隐式特征提取（TCN）：使用扩张因果卷积捕获多时间尺度的隐式依赖 + 2. 显式特征提取（S4）：使用结构化状态空间模型捕获长期显式依赖 + 3. 特征融合：将隐式和显式特征有机结合指导扩散去噪过程 + +模型架构： + 输入 -> DETACH分解 -> [隐式提取TCN, 显式提取S4] -> 特征融合 + -> 残差卷积层 + 扩散嵌入 -> 输出 + +超参数调节指南： + 1. 基础超参数： + - in_channels: 输入通道数（传感器/特征数量） + - res_channels: 残差层通道数 [建议: 128-512] + - skip_channels: 跳跃连接通道数 [建议: 128-512] + - num_res_layers: 残差层数量 [建议: 20-40] + - dropout: Dropout比率 [建议: 0.0-0.2] + + 2. 隐式模块调节（TCN）： + - tcn_channels: TCN通道数列表 [256, 256, 256] + - tcn_kernel_size: 卷积核大小 [3] + - tcn_dilation_rates: 扩张率 [1, 2, 4, 8] + * 调整扩张率可以改变时间尺度的捕获 + * 增加扩张率：[1, 2, 4, 8, 16] 捕获更长依赖 + * 跳跃增长：[1, 4, 16] 捕获稀疏长期依赖 + + 3. 显式模块调节（S4）： + - s4_d_state: S4状态维度 [建议: 32-128] + - s4_n_layers: S4层数 [建议: 2-6] + - s4_bidirectional: 双向模式 [True for imputation] + - S4的ABCD矩阵通过d_state和初始化方式控制 + + 4. 扩散模型参数： + - diffusion_step_embed_dim: 扩散步骤嵌入维度 + * 用于编码当前处于扩散过程的哪一步 + * 增大可以提供更精细的步骤信息 [建议: 128-512] + + 5. 训练技巧： + - 如果训练不稳定：减小学习率，增加warm-up步数 + - 如果效果不好：增大模型容量（res_channels, num_res_layers） + - 如果过拟合：增大dropout，使用数据增强 + - 如果计算太慢：减小num_res_layers，使用单向S4 +""" + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 导入TCN和S4模块 +from .tcn_implicit import TCNImplicitExtractor +from .s4_explicit import S4ExplicitExtractor + + +def swish(x): + """ + Swish激活函数: f(x) = x * sigmoid(x) + + 特点： + - 平滑的非线性函数 + - 自门控（self-gating）机制 + - 通常比ReLU效果更好 + """ + return x * torch.sigmoid(x) + + +def calc_diffusion_step_embedding(diffusion_steps, diffusion_step_embed_dim_in): + """ + 计算扩散步骤的位置编码 (Diffusion Step Embedding) + + 功能： + 将扩散步骤数（整数）编码为高维向量，类似于Transformer的位置编码 + + 原理： + 使用正弦和余弦函数的组合，为每个扩散步骤生成唯一的表示 + 这使模型能够识别当前处于扩散过程的哪个阶段 + + 参数： + diffusion_steps (torch.Tensor): 扩散步骤 [batch_size] 或 [batch_size, 1] + diffusion_step_embed_dim_in (int): 嵌入维度 + + 返回： + embed (torch.Tensor): 扩散步骤嵌入 [batch_size, diffusion_step_embed_dim_in] + + 数学公式： + PE(step, 2i) = sin(step / 10000^(2i/d)) + PE(step, 2i+1) = cos(step / 10000^(2i/d)) + """ + assert diffusion_step_embed_dim_in % 2 == 0, "嵌入维度必须是偶数" + + # 确保diffusion_steps是1维张量 + if len(diffusion_steps.shape) == 0: + diffusion_steps = diffusion_steps.unsqueeze(0) + elif len(diffusion_steps.shape) == 2: + diffusion_steps = diffusion_steps.squeeze(-1) + + half_dim = diffusion_step_embed_dim_in // 2 + + # 计算频率：10000^(-2i/d) + # 这使得低维度变化快，高维度变化慢，捕获不同频率的信息 + _embed = torch.log(torch.tensor(10000.0)) / (half_dim - 1) + _embed = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=diffusion_steps.device) * -_embed) + + # 计算 step * 频率 + _embed = diffusion_steps.float()[:, None] * _embed[None, :] # [B, half_dim] + + # 拼接sin和cos + embed = torch.cat([torch.sin(_embed), torch.cos(_embed)], dim=1) # [B, dim] + + return embed + + +class Conv1d(nn.Module): + """ + 带权重归一化的1D卷积层 + + 功能： + 标准的1D卷积，添加了权重归一化以提升训练稳定性 + """ + def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1): + super(Conv1d, self).__init__() + + # 计算padding以保持序列长度 + self.padding = dilation * (kernel_size - 1) // 2 + + # 1D卷积 + self.conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + dilation=dilation, + padding=self.padding + ) + + # 权重归一化：提升训练稳定性 + self.conv = nn.utils.weight_norm(self.conv) + + # Kaiming初始化（He初始化） + nn.init.kaiming_normal_(self.conv.weight) + + def forward(self, x): + return self.conv(x) + + +class ZeroConv1d(nn.Module): + """ + 零初始化的1D卷积层 + + 功能： + 权重和偏置都初始化为0的卷积层 + + 用途： + 在扩散模型中，零初始化的输出层可以使训练初期的预测接近原始输入 + 有助于稳定训练过程 + """ + def __init__(self, in_channels, out_channels): + super(ZeroConv1d, self).__init__() + + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=1, padding=0) + + # 零初始化 + self.conv.weight.data.zero_() + self.conv.bias.data.zero_() + + def forward(self, x): + return self.conv(x) + + +class ResidualBlock(nn.Module): + """ + 残差块 (Residual Block for Diffusion Model) + + 功能： + 扩散模型的核心构建块，整合了： + 1. 扩散步骤嵌入 + 2. 隐式和显式特征条件信息 + 3. 扩张卷积 + 4. 残差连接 + + 结构： + 输入 -> [+扩散嵌入] -> 扩张卷积 -> [+条件信息] -> 门控 -> [残差连接] -> 输出 + + 参数： + res_channels (int): 残差层通道数 + skip_channels (int): 跳跃连接通道数 + dilation (int): 扩张率 + diffusion_step_embed_dim_out (int): 扩散步骤嵌入输出维度 + cond_channels (int): 条件信息通道数 + """ + def __init__( + self, + res_channels, + skip_channels, + dilation, + diffusion_step_embed_dim_out, + cond_channels + ): + super(ResidualBlock, self).__init__() + + self.res_channels = res_channels + + # === 扩散步骤嵌入的全连接层 === + # 将扩散步骤嵌入映射到残差通道维度 + self.fc_diffusion = nn.Linear(diffusion_step_embed_dim_out, res_channels) + + # === 扩张卷积层 === + # 使用扩张卷积扩大感受野 + # 输出通道数是输入的2倍，用于门控机制（tanh和sigmoid） + self.dilated_conv = Conv1d( + res_channels, + 2 * res_channels, + kernel_size=3, + dilation=dilation + ) + + # === 条件信息卷积层 === + # 将条件信息（隐式+显式特征）映射到适合门控的维度 + self.cond_conv = Conv1d(cond_channels, 2 * res_channels, kernel_size=1) + + # === 残差连接卷积层 === + # 1x1卷积用于调整残差连接 + self.res_conv = nn.Conv1d(res_channels, res_channels, kernel_size=1) + self.res_conv = nn.utils.weight_norm(self.res_conv) + nn.init.kaiming_normal_(self.res_conv.weight) + + # === 跳跃连接卷积层 === + # 1x1卷积用于跳跃连接 + self.skip_conv = nn.Conv1d(res_channels, skip_channels, kernel_size=1) + self.skip_conv = nn.utils.weight_norm(self.skip_conv) + nn.init.kaiming_normal_(self.skip_conv.weight) + + def forward(self, x, cond, diffusion_step_embed): + """ + 前向传播 + + 参数： + x (torch.Tensor): 输入 [B, res_channels, L] + cond (torch.Tensor): 条件信息 [B, cond_channels, L] + diffusion_step_embed (torch.Tensor): 扩散步骤嵌入 [B, embed_dim] + + 返回： + residual_output (torch.Tensor): 残差输出 [B, res_channels, L] + skip_output (torch.Tensor): 跳跃连接输出 [B, skip_channels, L] + """ + B, C, L = x.shape + assert C == self.res_channels + + # === 1. 添加扩散步骤嵌入 === + # 将嵌入映射到残差通道维度，然后广播到序列长度 + diffusion_embed = self.fc_diffusion(diffusion_step_embed) # [B, res_channels] + diffusion_embed = diffusion_embed.view(B, self.res_channels, 1) # [B, res_channels, 1] + h = x + diffusion_embed # [B, res_channels, L] + + # === 2. 扩张卷积 === + h = self.dilated_conv(h) # [B, 2*res_channels, L] + + # === 3. 添加条件信息 === + assert cond is not None, "条件信息不能为None" + cond_out = self.cond_conv(cond) # [B, 2*res_channels, L] + h = h + cond_out # [B, 2*res_channels, L] + + # === 4. 门控机制 (Gated Activation) === + # 将通道分为两半，分别应用tanh和sigmoid，然后相乘 + # 这种机制来自WaveNet，可以有效地控制信息流 + h_tanh = torch.tanh(h[:, :self.res_channels, :]) # [B, res_channels, L] + h_sigmoid = torch.sigmoid(h[:, self.res_channels:, :]) # [B, res_channels, L] + h = h_tanh * h_sigmoid # [B, res_channels, L] + + # === 5. 残差连接和跳跃连接 === + residual_out = self.res_conv(h) # [B, res_channels, L] + skip_out = self.skip_conv(h) # [B, skip_channels, L] + + # 残差连接：输出 = (输入 + 残差) * sqrt(0.5) + # sqrt(0.5) 是缩放因子，用于训练稳定性 + residual_output = (x + residual_out) * math.sqrt(0.5) + + return residual_output, skip_out + + +class IEDiffusionModel(nn.Module): + """ + 隐式显式扩散模型 (Implicit-Explicit Diffusion Model) + + 功能： + 完整的扩散模型，整合TCN隐式特征和S4显式特征，用于时间序列插补 + + 架构流程： + 1. DETACH: 输入分解为 (noise, observed_data, mask) + 2. 隐式提取: TCN处理输入，提取多时间尺度隐式特征 + 3. 显式提取: S4处理输入，提取长期依赖显式特征 + 4. 特征融合: 合并隐式和显式特征 + 5. 条件准备: 将观测数据、mask和特征拼接作为条件 + 6. 残差处理: 通过多个残差块处理，每块结合扩散步骤嵌入 + 7. 输出生成: 聚合跳跃连接，生成去噪输出 + + 参数说明： + in_channels (int): 输入通道数（传感器/特征数量） + res_channels (int): 残差层通道数 + skip_channels (int): 跳跃连接通道数 + out_channels (int): 输出通道数（通常等于in_channels） + num_res_layers (int): 残差层数量 + dilation_cycle (int): 扩张率循环周期 + + diffusion_step_embed_dim_in (int): 扩散步骤嵌入输入维度 + diffusion_step_embed_dim_mid (int): 扩散步骤嵌入中间维度 + diffusion_step_embed_dim_out (int): 扩散步骤嵌入输出维度 + + tcn_channels (list): TCN各层通道数 + tcn_kernel_size (int): TCN卷积核大小 + tcn_dilation_rates (list): TCN扩张率列表 + tcn_dropout (float): TCN的dropout + + s4_d_state (int): S4状态维度 + s4_n_layers (int): S4层数 + s4_l_max (int): S4最大序列长度 + s4_dropout (float): S4的dropout + s4_bidirectional (bool): S4是否双向 + + 使用示例： + ```python + model = IEDiffusionModel( + in_channels=370, + res_channels=256, + skip_channels=256, + out_channels=370, + num_res_layers=36, + tcn_dilation_rates=[1, 2, 4, 8], + s4_d_state=64, + s4_bidirectional=True + ) + ``` + """ + def __init__( + self, + in_channels, + res_channels=256, + skip_channels=256, + out_channels=None, + num_res_layers=36, + dilation_cycle=10, + # 扩散步骤嵌入参数 + diffusion_step_embed_dim_in=128, + diffusion_step_embed_dim_mid=512, + diffusion_step_embed_dim_out=512, + # TCN隐式特征提取参数 + tcn_channels=[256, 256, 256], + tcn_kernel_size=3, + tcn_dilation_rates=[1, 2, 4, 8], + tcn_dropout=0.0, + # S4显式特征提取参数 + s4_d_state=64, + s4_n_layers=4, + s4_l_max=168, + s4_dropout=0.0, + s4_bidirectional=True + ): + super(IEDiffusionModel, self).__init__() + + # 如果未指定输出通道数，默认等于输入通道数 + if out_channels is None: + out_channels = in_channels + + self.in_channels = in_channels + self.res_channels = res_channels + self.skip_channels = skip_channels + self.out_channels = out_channels + self.num_res_layers = num_res_layers + + # ==================================================================== + # 1. 隐式特征提取模块 (TCN) + # ==================================================================== + self.tcn_extractor = TCNImplicitExtractor( + in_channels=in_channels, + hidden_channels=tcn_channels, + kernel_size=tcn_kernel_size, + dilation_rates=tcn_dilation_rates, + dropout=tcn_dropout + ) + + # ==================================================================== + # 2. 显式特征提取模块 (S4) + # ==================================================================== + # S4需要特征维度作为输入，这里使用res_channels + # 首先需要一个映射层将in_channels映射到res_channels + self.input_projection = nn.Conv1d(in_channels, res_channels, kernel_size=1) + + self.s4_extractor = S4ExplicitExtractor( + d_model=res_channels, + d_state=s4_d_state, + n_layers=s4_n_layers, + l_max=s4_l_max, + dropout=s4_dropout, + bidirectional=s4_bidirectional + ) + + # ==================================================================== + # 3. 初始卷积层 + # ==================================================================== + # 将噪声输入映射到残差通道维度 + self.init_conv = nn.Sequential( + Conv1d(in_channels, res_channels, kernel_size=1), + nn.ReLU() + ) + + # ==================================================================== + # 4. 扩散步骤嵌入网络 + # ==================================================================== + # 将扩散步骤编码映射到高维表示 + self.fc_diffusion1 = nn.Linear( + diffusion_step_embed_dim_in, + diffusion_step_embed_dim_mid + ) + self.fc_diffusion2 = nn.Linear( + diffusion_step_embed_dim_mid, + diffusion_step_embed_dim_out + ) + + # ==================================================================== + # 5. 计算条件信息通道数 + # ==================================================================== + # 条件信息包括：observed_data (in_channels) + mask (in_channels) + # + implicit_features (1) + explicit_features (res_channels) + cond_channels = in_channels + in_channels + 1 + res_channels + + # ==================================================================== + # 6. 残差块列表 + # ==================================================================== + self.residual_blocks = nn.ModuleList() + for n in range(num_res_layers): + # 计算扩张率：循环使用 1, 2, 4, ..., 2^(dilation_cycle-1) + dilation = 2 ** (n % dilation_cycle) + + block = ResidualBlock( + res_channels=res_channels, + skip_channels=skip_channels, + dilation=dilation, + diffusion_step_embed_dim_out=diffusion_step_embed_dim_out, + cond_channels=cond_channels + ) + self.residual_blocks.append(block) + + # ==================================================================== + # 7. 输出层 + # ==================================================================== + # 将跳跃连接聚合后映射到输出 + self.final_conv = nn.Sequential( + Conv1d(skip_channels, skip_channels, kernel_size=1), + nn.ReLU(), + ZeroConv1d(skip_channels, out_channels) + ) + + def forward(self, noise, observed_data, mask, diffusion_steps): + """ + 前向传播 + + 参数： + noise (torch.Tensor): 噪声输入 [B, in_channels, L] + observed_data (torch.Tensor): 观测数据 [B, in_channels, L] + mask (torch.Tensor): mask [B, in_channels, L] + 1表示观测值，0表示缺失值 + diffusion_steps (torch.Tensor): 扩散步骤 [B] 或 [B, 1] + + 返回： + output (torch.Tensor): 去噪输出 [B, out_channels, L] + + 处理流程： + 1. 提取隐式特征（TCN） + 2. 提取显式特征（S4） + 3. 准备条件信息 + 4. 初始化噪声输入 + 5. 计算扩散步骤嵌入 + 6. 通过残差块处理 + 7. 聚合跳跃连接 + 8. 生成最终输出 + """ + B, C, L = noise.shape + + # ==================================================================== + # 步骤1: 准备输入 - 将观测数据应用mask + # ==================================================================== + # 只保留观测值，缺失位置为0 + masked_data = observed_data * mask # [B, C, L] + + # ==================================================================== + # 步骤2: 提取隐式特征 (TCN) + # ==================================================================== + # TCN从masked_data中提取多时间尺度的隐式特征 + implicit_features = self.tcn_extractor(masked_data) # [B, 1, L] + + # ==================================================================== + # 步骤3: 提取显式特征 (S4) + # ==================================================================== + # 首先映射到res_channels维度 + s4_input = self.input_projection(masked_data) # [B, res_channels, L] + # S4提取长期依赖的显式特征 + explicit_features = self.s4_extractor(s4_input) # [B, res_channels, L] + + # ==================================================================== + # 步骤4: 准备条件信息 + # ==================================================================== + # 将所有条件信息拼接：observed_data, mask, implicit_features, explicit_features + conditional = torch.cat([ + masked_data, # [B, in_channels, L] + mask.float(), # [B, in_channels, L] + implicit_features, # [B, 1, L] + explicit_features # [B, res_channels, L] + ], dim=1) # [B, cond_channels, L] + + # ==================================================================== + # 步骤5: 初始化噪声输入 + # ==================================================================== + x = self.init_conv(noise) # [B, res_channels, L] + + # ==================================================================== + # 步骤6: 计算扩散步骤嵌入 + # ==================================================================== + # 将扩散步骤编码为高维向量 + diffusion_embed = calc_diffusion_step_embedding( + diffusion_steps, + self.fc_diffusion1.in_features + ) # [B, diffusion_step_embed_dim_in] + + # 通过全连接层映射 + diffusion_embed = swish(self.fc_diffusion1(diffusion_embed)) # [B, mid_dim] + diffusion_embed = swish(self.fc_diffusion2(diffusion_embed)) # [B, out_dim] + + # ==================================================================== + # 步骤7: 通过残差块处理 + # ==================================================================== + skip_sum = 0 # 累加所有跳跃连接 + + for block in self.residual_blocks: + # 每个残差块处理当前表示，并输出跳跃连接 + x, skip = block(x, conditional, diffusion_embed) + skip_sum = skip_sum + skip + + # 归一化跳跃连接：除以残差块数量的平方根，用于训练稳定性 + skip_sum = skip_sum / math.sqrt(self.num_res_layers) + + # ==================================================================== + # 步骤8: 生成最终输出 + # ==================================================================== + output = self.final_conv(skip_sum) # [B, out_channels, L] + + return output + + def get_config(self): + """ + 获取模型配置信息 + + 返回： + config (dict): 包含模型配置的字典 + """ + return { + 'in_channels': self.in_channels, + 'res_channels': self.res_channels, + 'skip_channels': self.skip_channels, + 'out_channels': self.out_channels, + 'num_res_layers': self.num_res_layers, + 'tcn_config': self.tcn_extractor.get_config(), + 's4_config': self.s4_extractor.get_config(), + 'num_parameters': sum(p.numel() for p in self.parameters()), + 'num_trainable_parameters': sum( + p.numel() for p in self.parameters() if p.requires_grad + ) + } + + +# ============================================================================ +# 使用示例和调试代码 +# ============================================================================ + +if __name__ == "__main__": + """ + 测试隐式显式扩散模型 + """ + print("=" * 80) + print("隐式显式扩散模型测试") + print("=" * 80) + + # 设置随机种子以便复现 + torch.manual_seed(42) + + # 模拟数据 + batch_size = 4 + in_channels = 370 # 370个传感器 + seq_len = 168 # 168小时（1周） + + # 创建随机数据 + noise = torch.randn(batch_size, in_channels, seq_len) + observed_data = torch.randn(batch_size, in_channels, seq_len) + + # 创建mask：30%缺失 + mask = torch.rand(batch_size, in_channels, seq_len) > 0.3 + mask = mask.float() + + # 扩散步骤：随机选择扩散步骤（0-199） + diffusion_steps = torch.randint(0, 200, (batch_size,)) + + print(f"\n输入形状:") + print(f" noise: {noise.shape}") + print(f" observed_data: {observed_data.shape}") + print(f" mask: {mask.shape} (缺失率: {1 - mask.mean():.2%})") + print(f" diffusion_steps: {diffusion_steps.shape}") + + # 创建模型 + model = IEDiffusionModel( + in_channels=in_channels, + res_channels=256, + skip_channels=256, + out_channels=in_channels, + num_res_layers=36, + # TCN参数 + tcn_channels=[256, 256, 256], + tcn_dilation_rates=[1, 2, 4, 8], + # S4参数 + s4_d_state=64, + s4_n_layers=4, + s4_l_max=seq_len, + s4_bidirectional=True + ) + + print(f"\n模型创建完成！") + + # 前向传播 + print(f"\n执行前向传播...") + output = model(noise, observed_data, mask, diffusion_steps) + print(f"输出形状: {output.shape}") + + # 显示模型配置 + config = model.get_config() + print(f"\n模型配置:") + print(f" 输入通道数: {config['in_channels']}") + print(f" 残差通道数: {config['res_channels']}") + print(f" 残差层数: {config['num_res_layers']}") + print(f" TCN感受野: {config['tcn_config']['receptive_field']} 时间步") + print(f" S4状态维度: {config['s4_config']['d_state']}") + print(f" S4层数: {config['s4_config']['n_layers']}") + print(f" 参数总数: {config['num_parameters']:,}") + print(f" 可训练参数: {config['num_trainable_parameters']:,}") + + print("\n" + "=" * 80) + print("测试完成！") + print("=" * 80) diff --git a/IEclaude/models/s4_explicit.py b/IEclaude/models/s4_explicit.py new file mode 100644 index 0000000..3c1211c --- /dev/null +++ b/IEclaude/models/s4_explicit.py @@ -0,0 +1,550 @@ +""" +S4显式特征提取模块 (Structured State Space Model for Explicit Feature Extraction) + +功能说明： + 使用结构化状态空间模型提取时间序列的长期依赖特征 + +核心概念： + S4是一种序列建模方法，通过状态空间模型捕获长期依赖关系 + + 状态空间模型的离散形式： + x_{t+1} = A * x_t + B * u_t (状态更新方程) + y_t = C * x_t + D * u_t (输出方程) + + 其中： + - u_t: 输入序列 [input sequence] + - x_t: 隐藏状态 [hidden state] + - y_t: 输出序列 [output sequence] + - A: 状态转移矩阵 [NxN] - 控制历史状态如何演化 + - B: 输入矩阵 [Nx1] - 控制当前输入如何影响状态 + - C: 输出矩阵 [1xN] - 控制状态如何映射到输出 + - D: 前馈矩阵 [标量] - 控制输入到输出的直接连接（skip connection） + +超参数调节指南： + 1. 基础超参数： + - d_model: 特征维度，即隐藏层大小 [建议: 64-512] + - d_state (N): 状态维度，控制模型容量 [建议: 16-128] + * 增大N可以提升模型对复杂序列的建模能力 + * 但会增加计算量，建议从64开始 + - dropout: 防止过拟合 [建议: 0.0-0.2] + + 2. 架构调整 - ABCD矩阵： + - A矩阵初始化方式： + * 'hippo': HiPPO初始化，适合长期依赖（推荐） + * 'diagonal': 对角化初始化，计算更快但表达能力较弱 + * 'random': 随机初始化 + + - B, C矩阵： + * 一般使用标准正态分布初始化 + * 可以通过学习率控制这些参数的学习速度 + + - D矩阵： + * 相当于残差连接的权重 + * 通常初始化为较小的值（0.1-1.0） + + 3. 序列长度： + - l_max: 最大序列长度 [建议: 设置为实际序列长度] + * 对于交通数据：如果用168小时（1周），设为168 + * 如果不确定，可以设为2的幂次（如256, 512） + + 4. 双向模式： + - bidirectional: 是否使用双向S4 [True/False] + * True: 同时考虑过去和未来的信息（适合插补任务） + * False: 只考虑过去的信息（适合预测任务） + +调优建议： + - 如果训练不稳定：减小学习率，增大d_state + - 如果效果不好：尝试双向模式，增大d_state + - 如果过拟合：增大dropout，减小d_state + - 如果计算太慢：减小d_state，使用单向模式 +""" + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy import special as ss +import numpy as np + + +def hippo_initializer(N): + """ + HiPPO (High-order Polynomial Projection Operators) 初始化 + + 功能： + 生成适合建模长期依赖的A和B矩阵初始值 + + 原理： + HiPPO理论证明，通过特定的A、B矩阵设计，可以使状态向量 + 高效地记忆和压缩历史信息，实现对长序列的有效建模 + + 参数： + N (int): 状态维度 + + 返回： + A (np.array): 状态转移矩阵 [N, N] + B (np.array): 输入矩阵 [N, 1] + + 数学推导： + HiPPO矩阵的构造基于Legendre多项式，能够将历史信息 + 投影到正交多项式基上，实现信息的有效压缩和记忆 + """ + # HiPPO-LegS 矩阵构造 + # 这是HiPPO论文中推荐的矩阵形式，特别适合长期依赖建模 + + # 创建NxN的下三角矩阵 + # A矩阵的(n,k)元素定义如下： + A = np.zeros((N, N)) + for n in range(N): + for k in range(N): + if n > k: + A[n, k] = np.sqrt(2 * n + 1) * np.sqrt(2 * k + 1) + elif n == k: + A[n, k] = n + 1 + + # B矩阵定义为每个元素 B[n] = sqrt(2*n + 1) + B = np.sqrt(2 * np.arange(N) + 1).reshape(N, 1) + + return A, B + + +def discretize_zoh(A, B, dt): + """ + 零阶保持器离散化 (Zero-Order Hold Discretization) + + 功能： + 将连续状态空间模型离散化为离散状态空间模型 + + 原理： + 连续系统: dx/dt = A*x + B*u + 离散系统: x_{k+1} = A_bar*x_k + B_bar*u_k + + 零阶保持假设输入信号在采样间隔内保持恒定 + + 参数： + A (torch.Tensor): 连续状态转移矩阵 [N, N] + B (torch.Tensor): 连续输入矩阵 [N, 1] + dt (float): 采样时间步长（离散化步长） + + 返回： + A_bar (torch.Tensor): 离散状态转移矩阵 [N, N] + B_bar (torch.Tensor): 离散输入矩阵 [N, 1] + + 数学公式： + A_bar = exp(A * dt) + B_bar = (A^{-1})(exp(A*dt) - I) * B + + 调节建议： + - dt越小，离散化越精确，但可能需要更多时间步 + - dt越大，计算越快，但可能损失精度 + - 一般设置dt=1.0即可 + """ + # 计算矩阵指数 exp(A * dt) + # 这里使用特征值分解来高效计算 + I = torch.eye(A.shape[0], device=A.device, dtype=A.dtype) + A_bar = torch.matrix_exp(A * dt) + + # 计算 B_bar = (A^{-1})(exp(A*dt) - I) * B + # 如果A可逆，使用该公式；否则使用近似 + try: + A_inv = torch.linalg.inv(A) + B_bar = A_inv @ (A_bar - I) @ B + except: + # 如果A不可逆，使用一阶近似：B_bar ≈ dt * B + B_bar = dt * B + + return A_bar, B_bar + + +class S4Kernel(nn.Module): + """ + S4卷积核模块 (S4 Convolutional Kernel) + + 功能： + 生成S4的卷积核，用于高效处理长序列 + + 原理： + S4通过将状态空间模型转换为卷积形式，可以利用FFT加速计算 + 卷积核K的每个元素: K_l = C * A^l * B + 这样，输出可以表示为: y = K * u (卷积操作) + + 参数： + d_model (int): 特征维度 + d_state (int): 状态维度 N + l_max (int): 最大序列长度 + dt_min (float): 最小时间步长 + dt_max (float): 最大时间步长 + + ABCD矩阵调节： + - 通过修改A的初始化方式改变记忆模式 + - 通过修改B、C的初始化改变输入输出的映射 + - 通过log_dt调节时间尺度 + """ + def __init__(self, d_model, d_state=64, l_max=1, dt_min=0.001, dt_max=0.1): + super().__init__() + + self.d_model = d_model # H：特征维度 + self.d_state = d_state # N：状态维度 + self.l_max = l_max # L：最大序列长度 + + # === A矩阵：状态转移矩阵 [N, N] === + # 使用HiPPO初始化，适合长期依赖 + A, B_init = hippo_initializer(d_state) + self.A = nn.Parameter(torch.tensor(A, dtype=torch.float32)) + + # === B矩阵：输入矩阵 [N, 1] === + # 控制当前输入如何影响状态 + # 使用HiPPO推荐的初始化 + self.B = nn.Parameter(torch.tensor(B_init, dtype=torch.float32)) + + # === C矩阵：输出矩阵 [1, N] === + # 控制状态如何映射到输出 + # 使用标准正态分布初始化 + C_init = torch.randn(1, d_state) / np.sqrt(d_state) + self.C = nn.Parameter(C_init) + + # === log_dt: 对数时间步长 === + # 通过学习时间步长，模型可以自适应调整时间尺度 + # dt越大，状态更新越激进；dt越小，状态更新越保守 + log_dt = torch.rand(d_model) * ( + np.log(dt_max) - np.log(dt_min) + ) + np.log(dt_min) + self.log_dt = nn.Parameter(log_dt) + + # === D矩阵：前馈矩阵（skip connection） === + # 允许输入直接影响输出，类似于残差连接 + # 初始化为1.0，可以调节以控制skip connection的强度 + self.D = nn.Parameter(torch.ones(d_model)) + + def kernel(self, L): + """ + 生成长度为L的S4卷积核 + + 参数： + L (int): 目标序列长度 + + 返回： + K (torch.Tensor): 卷积核 [d_model, L] + + 计算流程： + 1. 离散化A、B矩阵 + 2. 计算K_l = C * A^l * B for l=0,1,...,L-1 + 3. 返回卷积核 + """ + # 计算时间步长 dt = exp(log_dt) + dt = torch.exp(self.log_dt) # [d_model] + + # 为每个特征维度离散化状态空间 + # 这里简化处理：使用平均dt进行离散化 + dt_mean = dt.mean() + + # 离散化A和B矩阵 + A_bar, B_bar = discretize_zoh(self.A, self.B, dt_mean) + + # 计算卷积核 K_l = C * (A_bar)^l * B_bar + # 使用幂迭代方法计算 + K = [] + A_power = torch.eye(self.d_state, device=self.A.device, dtype=self.A.dtype) + + for l in range(L): + # K_l = C * A^l * B + K_l = self.C @ A_power @ B_bar # [1, N] @ [N, N] @ [N, 1] = [1, 1] + K.append(K_l.squeeze()) + + # 更新 A^l 为 A^{l+1} + A_power = A_power @ A_bar + + # 将列表转换为张量 [L] + K = torch.stack(K, dim=0) # [L] + + # 扩展到所有特征维度 [d_model, L] + K = K.unsqueeze(0).expand(self.d_model, -1) + + return K + + def forward(self, L): + """ + 前向传播：生成卷积核 + + 参数： + L (int): 序列长度 + + 返回： + K (torch.Tensor): 卷积核 [d_model, L] + """ + return self.kernel(L) + + +class S4Layer(nn.Module): + """ + S4层 (S4 Layer) + + 功能： + 完整的S4层，包括S4卷积、激活函数、dropout和输出线性变换 + + 结构： + 输入 -> S4卷积 -> 激活 -> Dropout -> 线性变换 -> 输出 + + 参数： + d_model (int): 特征维度 + d_state (int): 状态维度N + l_max (int): 最大序列长度 + dropout (float): Dropout比率 + bidirectional (bool): 是否使用双向S4 + activation (str): 激活函数类型 ['gelu', 'relu', 'swish'] + + 使用建议： + - 对于插补任务，建议使用bidirectional=True + - 对于预测任务，建议使用bidirectional=False + - 激活函数通常选择'gelu'，效果较好 + """ + def __init__( + self, + d_model, + d_state=64, + l_max=1, + dropout=0.0, + bidirectional=True, + activation='gelu' + ): + super().__init__() + + self.d_model = d_model + self.d_state = d_state + self.bidirectional = bidirectional + + # S4卷积核 + self.kernel = S4Kernel(d_model, d_state, l_max) + + # 如果使用双向，需要两个kernel（前向和后向） + if bidirectional: + self.kernel_backward = S4Kernel(d_model, d_state, l_max) + + # 激活函数 + if activation == 'gelu': + self.activation = nn.GELU() + elif activation == 'relu': + self.activation = nn.ReLU() + elif activation == 'swish': + self.activation = nn.SiLU() # Swish = SiLU + else: + self.activation = nn.Identity() + + # Dropout + self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() + + # 输出线性变换 + # 如果是双向，特征维度会翻倍 + output_dim = d_model * 2 if bidirectional else d_model + self.output_linear = nn.Linear(output_dim, d_model) + + # Layer Normalization：稳定训练 + self.norm = nn.LayerNorm(d_model) + + def forward(self, u): + """ + 前向传播 + + 参数： + u (torch.Tensor): 输入 [batch_size, d_model, seq_len] + + 返回： + y (torch.Tensor): 输出 [batch_size, d_model, seq_len] + + 计算流程： + 1. 生成S4卷积核 + 2. 使用FFT进行高效卷积 + 3. 添加skip connection (D矩阵) + 4. 激活、dropout和输出变换 + """ + B, H, L = u.shape + assert H == self.d_model + + # === 生成卷积核 === + k = self.kernel(L) # [d_model, L] + + # === 前向卷积 === + # 使用FFT加速卷积运算 + # 卷积定理: conv(x, k) = IFFT(FFT(x) * FFT(k)) + k_f = torch.fft.rfft(k, n=2*L) # [d_model, L_fft] + u_f = torch.fft.rfft(u, n=2*L, dim=-1) # [B, d_model, L_fft] + y_f = u_f * k_f.unsqueeze(0) # [B, d_model, L_fft] + y = torch.fft.irfft(y_f, n=2*L)[..., :L] # [B, d_model, L] + + # === 双向处理 === + if self.bidirectional: + # 后向卷积：反转序列 + k_b = self.kernel_backward(L) + k_b_f = torch.fft.rfft(k_b, n=2*L) + u_reversed = torch.flip(u, dims=[-1]) # 反转序列 + u_reversed_f = torch.fft.rfft(u_reversed, n=2*L, dim=-1) + y_b_f = u_reversed_f * k_b_f.unsqueeze(0) + y_b = torch.fft.irfft(y_b_f, n=2*L)[..., :L] + y_b = torch.flip(y_b, dims=[-1]) # 反转回来 + + # 拼接前向和后向 + y = torch.cat([y, y_b], dim=1) # [B, 2*d_model, L] + + # === 添加skip connection (D矩阵) === + D = self.kernel.D.unsqueeze(0).unsqueeze(-1) # [1, d_model, 1] + y = y + u * D # [B, d_model, L] 或 [B, 2*d_model, L] + + # === 激活和Dropout === + y = self.activation(y) + y = self.dropout(y) + + # === 输出线性变换 === + # 转换为 [B, L, d_model] 格式进行线性变换 + y = y.transpose(1, 2) # [B, L, d_model] 或 [B, L, 2*d_model] + y = self.output_linear(y) # [B, L, d_model] + + # === Layer Normalization + Residual Connection === + # 转换回 [B, d_model, L] + y = y.transpose(1, 2) # [B, d_model, L] + + # 残差连接 + u_norm = self.norm(u.transpose(1, 2)).transpose(1, 2) + y = y + u_norm + + return y + + +class S4ExplicitExtractor(nn.Module): + """ + S4显式特征提取器 (S4 Explicit Feature Extractor) + + 功能： + 堆叠多个S4层，提取长期依赖的显式特征 + + 架构： + 输入 -> [S4Layer1, S4Layer2, ..., S4LayerN] -> 输出 + + 参数： + d_model (int): 特征维度 + d_state (int): 状态维度N + n_layers (int): S4层数 + l_max (int): 最大序列长度 + dropout (float): Dropout比率 + bidirectional (bool): 是否使用双向S4 + + 调优建议： + 1. 增加n_layers可以提取更抽象的特征（2-6层） + 2. 增大d_state可以提升建模能力（32-128） + 3. 对于长序列（>100），建议使用bidirectional=True + 4. 如果过拟合，增大dropout（0.1-0.2） + """ + def __init__( + self, + d_model, + d_state=64, + n_layers=4, + l_max=1, + dropout=0.0, + bidirectional=True + ): + super().__init__() + + self.d_model = d_model + self.d_state = d_state + self.n_layers = n_layers + + # 堆叠多个S4层 + self.layers = nn.ModuleList([ + S4Layer( + d_model=d_model, + d_state=d_state, + l_max=l_max, + dropout=dropout, + bidirectional=bidirectional + ) + for _ in range(n_layers) + ]) + + # 最终的Layer Normalization + self.final_norm = nn.LayerNorm(d_model) + + def forward(self, x): + """ + 前向传播 + + 参数： + x (torch.Tensor): 输入 [batch_size, d_model, seq_len] + + 返回： + explicit_features (torch.Tensor): 显式特征 [batch_size, d_model, seq_len] + """ + # 依次通过所有S4层 + for layer in self.layers: + x = layer(x) + + # 最终归一化 + # 转换为 [B, L, d_model] 进行LayerNorm + x = x.transpose(1, 2) + x = self.final_norm(x) + x = x.transpose(1, 2) # 转换回 [B, d_model, L] + + return x + + def get_config(self): + """ + 获取模型配置信息 + + 返回： + config (dict): 包含模型配置的字典 + """ + return { + 'd_model': self.d_model, + 'd_state': self.d_state, + 'n_layers': self.n_layers, + 'num_parameters': sum(p.numel() for p in self.parameters()), + 'num_trainable_parameters': sum(p.numel() for p in self.parameters() if p.requires_grad) + } + + +# ============================================================================ +# 使用示例和调试代码 +# ============================================================================ + +if __name__ == "__main__": + """ + 测试S4显式特征提取器 + """ + print("=" * 80) + print("S4显式特征提取器测试") + print("=" * 80) + + # 模拟输入数据 + batch_size = 4 + d_model = 256 # 特征维度 + seq_len = 168 # 序列长度（1周） + + # 创建随机输入 + x = torch.randn(batch_size, d_model, seq_len) + print(f"\n输入张量形状: {x.shape}") + + # 创建S4模型 + model = S4ExplicitExtractor( + d_model=d_model, + d_state=64, # 状态维度 + n_layers=4, # 4层S4 + l_max=seq_len, # 最大序列长度 + dropout=0.1, + bidirectional=True # 双向 + ) + + # 前向传播 + output = model(x) + print(f"输出张量形状: {output.shape}") + + # 显示配置 + config = model.get_config() + print(f"\n模型配置:") + print(f" 特征维度: {config['d_model']}") + print(f" 状态维度: {config['d_state']}") + print(f" 层数: {config['n_layers']}") + print(f" 参数总数: {config['num_parameters']:,}") + print(f" 可训练参数: {config['num_trainable_parameters']:,}") + + print("\n" + "=" * 80) + print("测试完成！") + print("=" * 80) diff --git a/IEclaude/models/tcn_implicit.py b/IEclaude/models/tcn_implicit.py new file mode 100644 index 0000000..b506dd3 --- /dev/null +++ b/IEclaude/models/tcn_implicit.py @@ -0,0 +1,391 @@ +""" +TCN隐式特征提取模块 (Temporal Convolutional Network for Implicit Feature Extraction) + +功能说明： + 使用扩张因果卷积提取时间序列的隐式特征，捕获不同时间尺度的依赖关系 + +核心概念： + - 因果卷积：保证时间t的输出只依赖于t及之前的输入，不会产生信息泄露 + - 扩张卷积：通过调整扩张率(dilation rate)来扩大感受野，捕获更长时间尺度的依赖 + - 多尺度特征：不同扩张率的卷积层可以提取不同时间尺度的特征 + +超参数调节指南： + 1. 基础超参数： + - channels: 每层的通道数，增大可提升表达能力但会增加计算量 [建议: 64-256] + - kernel_size: 卷积核大小，影响局部感受野 [建议: 3] + - dropout: 防止过拟合的dropout比率 [建议: 0.0-0.2] + + 2. 架构调整： + - dilation_rates: 扩张率序列，控制时间尺度 + * [1, 2, 4, 8]: 标准指数增长，捕获1到8步的依赖 + * [1, 2, 4, 8, 16]: 更长的依赖关系 + * [1, 4, 16]: 跳跃式增长，捕获更稀疏的长期依赖 + * 计算感受野: RF = 1 + 2 * (kernel_size - 1) * sum(dilation_rates) + + - num_layers: TCN层数，增加深度可以提取更抽象的特征 [建议: 3-6] + + 3. 优化技巧： + - 使用weight normalization提升训练稳定性 + - 使用residual connection加速收敛 + - 使用dropout防止过拟合 +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class CausalConv1d(nn.Module): + """ + 因果卷积层 (Causal Convolution Layer) + + 特点： + - 保证因果性：输出只依赖于当前和过去的输入 + - 使用padding确保输出长度与输入相同 + + 参数说明： + in_channels (int): 输入通道数 + out_channels (int): 输出通道数 + kernel_size (int): 卷积核大小，通常设为3 + dilation (int): 扩张率，控制卷积核元素之间的间隔 + - dilation=1: 标准卷积，感受野为kernel_size + - dilation=2: 每隔一个位置采样，感受野扩大2倍 + - dilation=4: 每隔三个位置采样，感受野扩大4倍 + dropout (float): Dropout比率，用于正则化 + + 输入输出： + 输入: [batch_size, in_channels, seq_len] + 输出: [batch_size, out_channels, seq_len] + """ + def __init__(self, in_channels, out_channels, kernel_size, dilation, dropout=0.0): + super(CausalConv1d, self).__init__() + + # 计算因果卷积所需的padding + # padding = (kernel_size - 1) * dilation 确保输出长度等于输入长度 + # 这个padding会添加在序列的左侧（过去），保证因果性 + self.padding = (kernel_size - 1) * dilation + + # 一维卷积层 + # dilation参数控制卷积核元素之间的间隔，实现不同时间尺度的特征提取 + self.conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + padding=0, # 我们手动处理padding以实现因果性 + dilation=dilation + ) + + # Weight Normalization: 将权重向量分解为方向和幅度 + # 优点：加速收敛，提高训练稳定性，对初始化不敏感 + self.conv = nn.utils.weight_norm(self.conv) + + # Dropout: 训练时随机丢弃一些神经元，防止过拟合 + self.dropout = nn.Dropout(dropout) + + # 使用Kaiming初始化（He初始化） + # 适用于ReLU激活函数，有助于避免梯度消失/爆炸 + nn.init.kaiming_normal_(self.conv.weight) + + def forward(self, x): + """ + 前向传播 + + 参数： + x: 输入张量 [batch_size, in_channels, seq_len] + + 返回： + out: 输出张量 [batch_size, out_channels, seq_len] + """ + # 在序列左侧（过去）添加padding，保证因果性 + # F.pad的参数(self.padding, 0)表示在最后一维的左侧填充self.padding个0，右侧不填充 + x = F.pad(x, (self.padding, 0)) + + # 执行卷积操作 + x = self.conv(x) + + # 应用dropout + x = self.dropout(x) + + return x + + +class TCNResidualBlock(nn.Module): + """ + TCN残差块 (TCN Residual Block) + + 结构： + 输入 -> CausalConv1d -> ReLU -> CausalConv1d -> ReLU -> 输出 + | | + +-------------------residual connection----------------+ + + 优点： + - 残差连接允许梯度直接传播，解决深层网络训练困难的问题 + - 使网络能够学习残差（增量），而不是完整的变换 + - 允许堆叠更多层以提取更抽象的特征 + + 参数说明： + channels (int): 输入输出通道数（残差块保持通道数不变） + kernel_size (int): 卷积核大小 + dilation (int): 扩张率 + dropout (float): Dropout比率 + """ + def __init__(self, channels, kernel_size, dilation, dropout=0.0): + super(TCNResidualBlock, self).__init__() + + # 第一个因果卷积层 + # 使用相同的dilation rate，保持时间尺度一致 + self.conv1 = CausalConv1d( + channels, channels, kernel_size, dilation, dropout + ) + + # 第二个因果卷积层 + # 进一步提取特征，增加网络的非线性表达能力 + self.conv2 = CausalConv1d( + channels, channels, kernel_size, dilation, dropout + ) + + # ReLU激活函数：max(0, x) + # 引入非线性，使网络能够拟合复杂函数 + self.relu = nn.ReLU() + + def forward(self, x): + """ + 前向传播 + + 参数： + x: 输入张量 [batch_size, channels, seq_len] + + 返回： + out: 输出张量 [batch_size, channels, seq_len] + """ + # 保存输入，用于残差连接 + residual = x + + # 第一个卷积 -> 激活 + out = self.relu(self.conv1(x)) + + # 第二个卷积 -> 激活 + out = self.relu(self.conv2(out)) + + # 残差连接：输出 = F(x) + x + # 这使得网络学习残差F(x)而不是完整的映射H(x) + # 如果F(x)=0，则网络退化为恒等映射，不会降低性能 + out = out + residual + + return out + + +class TCNImplicitExtractor(nn.Module): + """ + TCN隐式特征提取器 (TCN Implicit Feature Extractor) + + 功能： + 通过堆叠多个不同扩张率的TCN残差块，提取多时间尺度的隐式特征 + + 架构设计： + 输入 -> 输入卷积 -> [TCN残差块1, TCN残差块2, ..., TCN残差块N] -> 输出卷积 -> 隐式特征 + + 多尺度特征提取原理： + - 第1层 (dilation=1): 捕获相邻时间步的局部特征 + - 第2层 (dilation=2): 捕获间隔1步的短期依赖 + - 第3层 (dilation=4): 捕获间隔3步的中期依赖 + - 第4层 (dilation=8): 捕获间隔7步的长期依赖 + + 参数说明： + in_channels (int): 输入通道数，对应时间序列的特征维度 + hidden_channels (list): 每层TCN的通道数列表 + - 例如 [64, 64, 64] 表示3层，每层64个通道 + - 可以设置为递增 [64, 128, 256] 以逐层提取更抽象的特征 + kernel_size (int): 卷积核大小，通常设为3 + dilation_rates (list): 扩张率列表，每个元素对应一个TCN层 + - [1, 2, 4, 8]: 指数增长，适合大多数时间序列 + - [1, 4, 16]: 跳跃增长，适合需要捕获稀疏长期依赖的情况 + dropout (float): Dropout比率 + + 调优建议： + 1. 如果模型过拟合：增大dropout (0.1 -> 0.2) + 2. 如果需要更长的依赖：增加dilation_rates (添加16, 32等) + 3. 如果特征不够丰富：增加hidden_channels (64 -> 128) + 4. 如果训练不稳定：减小学习率或增加batch normalization + """ + def __init__( + self, + in_channels, + hidden_channels=[256, 256, 256], + kernel_size=3, + dilation_rates=[1, 2, 4, 8], + dropout=0.0 + ): + super(TCNImplicitExtractor, self).__init__() + + # 输入卷积：将输入维度映射到隐藏维度 + # 使用1x1卷积（kernel_size=1）进行通道数变换，不改变序列长度 + self.input_conv = nn.Conv1d(in_channels, hidden_channels[0], kernel_size=1) + self.input_conv = nn.utils.weight_norm(self.input_conv) + nn.init.kaiming_normal_(self.input_conv.weight) + + # TCN残差块列表 + self.tcn_blocks = nn.ModuleList() + + # 为每个扩张率创建一个TCN残差块 + # 如果hidden_channels长度小于dilation_rates，则重复使用最后一个channel数 + for i, dilation in enumerate(dilation_rates): + # 确定当前层的通道数 + if i < len(hidden_channels): + channels = hidden_channels[i] + else: + channels = hidden_channels[-1] + + # 如果不是第一层，且通道数改变，需要添加通道变换层 + if i > 0 and channels != self.tcn_blocks[-1].conv1.conv.out_channels: + # 添加1x1卷积进行通道数变换 + channel_transform = nn.Conv1d( + self.tcn_blocks[-1].conv1.conv.out_channels, + channels, + kernel_size=1 + ) + channel_transform = nn.utils.weight_norm(channel_transform) + self.tcn_blocks.append(channel_transform) + + # 添加TCN残差块 + # 每个块使用不同的扩张率，捕获不同时间尺度的特征 + block = TCNResidualBlock( + channels, + kernel_size, + dilation, + dropout + ) + self.tcn_blocks.append(block) + + # 输出卷积：将隐藏维度映射到输出维度 + # 这里输出通道数设为1，用于生成标量形式的隐式特征 + final_channels = hidden_channels[-1] if hidden_channels else hidden_channels[0] + self.output_conv = nn.Conv1d(final_channels, 1, kernel_size=1) + self.output_conv = nn.utils.weight_norm(self.output_conv) + nn.init.kaiming_normal_(self.output_conv.weight) + + # 保存配置用于调试和模型分析 + self.in_channels = in_channels + self.hidden_channels = hidden_channels + self.dilation_rates = dilation_rates + + # 计算理论感受野（Receptive Field） + # 感受野表示输出神经元能"看到"的输入范围 + self.receptive_field = self._calculate_receptive_field(kernel_size, dilation_rates) + + def _calculate_receptive_field(self, kernel_size, dilation_rates): + """ + 计算TCN的理论感受野 + + 感受野计算公式： + RF = 1 + 2 * (kernel_size - 1) * sum(dilation_rates) + + 例如：kernel_size=3, dilation_rates=[1,2,4,8] + RF = 1 + 2 * (3-1) * (1+2+4+8) = 1 + 2*2*15 = 61 + + 这意味着输出的每个位置可以"看到"过去61个时间步的信息 + """ + rf = 1 + for dilation in dilation_rates: + rf += 2 * (kernel_size - 1) * dilation + return rf + + def forward(self, x): + """ + 前向传播 + + 参数： + x: 输入张量 [batch_size, in_channels, seq_len] + - batch_size: 批次大小 + - in_channels: 输入特征维度（例如交通传感器数量） + - seq_len: 序列长度（例如时间步数） + + 返回： + implicit_features: 隐式特征 [batch_size, 1, seq_len] + 或 [batch_size, hidden_channels[-1], seq_len] + + 处理流程： + 1. 输入卷积：调整通道数 + 2. 依次通过所有TCN残差块：提取多尺度特征 + 3. 输出卷积：生成最终的隐式特征表示 + """ + # 输入卷积：[B, in_channels, L] -> [B, hidden_channels[0], L] + out = self.input_conv(x) + + # 通过所有TCN残差块 + # 每个块提取特定时间尺度的特征 + for block in self.tcn_blocks: + out = block(out) + + # 输出卷积：生成隐式特征 + # [B, hidden_channels[-1], L] -> [B, 1, L] 或保持多通道 + implicit_features = self.output_conv(out) + + return implicit_features + + def get_config(self): + """ + 获取模型配置信息，用于模型分析和调试 + + 返回： + config (dict): 包含模型配置的字典 + """ + return { + 'in_channels': self.in_channels, + 'hidden_channels': self.hidden_channels, + 'dilation_rates': self.dilation_rates, + 'receptive_field': self.receptive_field, + 'num_parameters': sum(p.numel() for p in self.parameters()), + 'num_trainable_parameters': sum(p.numel() for p in self.parameters() if p.requires_grad) + } + + +# ============================================================================ +# 使用示例和调试代码 +# ============================================================================ + +if __name__ == "__main__": + """ + 测试TCN隐式特征提取器 + + 这个测试代码展示了如何使用TCN模块，以及如何查看模型配置 + """ + print("=" * 80) + print("TCN隐式特征提取器测试") + print("=" * 80) + + # 模拟输入数据 + batch_size = 4 # 批次大小 + in_channels = 370 # 输入通道数（例如370个交通传感器） + seq_len = 168 # 序列长度（例如168小时=1周） + + # 创建随机输入张量 + x = torch.randn(batch_size, in_channels, seq_len) + print(f"\n输入张量形状: {x.shape}") + + # 创建TCN模型 + model = TCNImplicitExtractor( + in_channels=in_channels, + hidden_channels=[256, 256, 256], # 3层，每层256通道 + kernel_size=3, + dilation_rates=[1, 2, 4, 8], # 4个不同的时间尺度 + dropout=0.1 + ) + + # 前向传播 + output = model(x) + print(f"输出张量形状: {output.shape}") + + # 显示模型配置 + config = model.get_config() + print(f"\n模型配置:") + print(f" 输入通道数: {config['in_channels']}") + print(f" 隐藏层通道数: {config['hidden_channels']}") + print(f" 扩张率: {config['dilation_rates']}") + print(f" 感受野: {config['receptive_field']} 时间步") + print(f" 参数总数: {config['num_parameters']:,}") + print(f" 可训练参数: {config['num_trainable_parameters']:,}") + + print("\n" + "=" * 80) + print("测试完成！") + print("=" * 80) diff --git a/IEclaude/requirements.txt b/IEclaude/requirements.txt new file mode 100644 index 0000000..da1f6c1 --- /dev/null +++ b/IEclaude/requirements.txt @@ -0,0 +1,27 @@ +# IEclaude Requirements +# Python 3.8+ + +# 深度学习框架 +torch>=1.12.0 +torchvision>=0.13.0 + +# 数据处理 +numpy>=1.21.0 +pandas>=1.3.0 +scipy>=1.7.0 + +# 机器学习工具 +scikit-learn>=1.0.0 + +# 可视化 +matplotlib>=3.4.0 +seaborn>=0.11.0 + +# 进度条 +tqdm>=4.62.0 + +# JSON处理 (Python内置，无需安装) +# json + +# 其他工具 +einops>=0.6.0 # 用于S4模块的张量操作 diff --git a/IEclaude/run_all.sh b/IEclaude/run_all.sh new file mode 100755 index 0000000..550de6d --- /dev/null +++ b/IEclaude/run_all.sh @@ -0,0 +1,169 @@ +#!/bin/bash +# ============================================================================ +# IEclaude 完整训练和评估脚本 +# +# 功能： +# 1. 生成所有配置文件 +# 2. 依次训练所有缺失率的模型 +# 3. 评估所有模型 +# 4. 生成汇总报告 +# +# 使用方法： +# chmod +x run_all.sh +# ./run_all.sh --gpu 0 +# ============================================================================ + +# 设置颜色输出 +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +RED='\033[0;31m' +NC='\033[0m' # No Color + +# 解析命令行参数 +GPU=0 +if [ "$1" == "--gpu" ]; then + GPU=$2 +fi + +echo -e "${GREEN}========================================${NC}" +echo -e "${GREEN}IEclaude 完整训练和评估流程${NC}" +echo -e "${GREEN}========================================${NC}" +echo -e "使用GPU: ${GPU}" +echo "" + +# 设置GPU环境变量 +export CUDA_VISIBLE_DEVICES=${GPU} + +# 缺失率列表 +MISSING_RATES=(20 30 40 50 60 70 80) + +# ============================================================================ +# 步骤1: 生成配置文件 +# ============================================================================ +echo -e "${YELLOW}步骤1: 生成配置文件${NC}" +python generate_configs.py +echo "" + +# ============================================================================ +# 步骤2: 训练所有模型 +# ============================================================================ +echo -e "${YELLOW}步骤2: 训练所有模型${NC}" + +for RATE in "${MISSING_RATES[@]}"; do + echo -e "${GREEN}========================================${NC}" + echo -e "${GREEN}训练缺失率 ${RATE}% 的模型${NC}" + echo -e "${GREEN}========================================${NC}" + + python train.py \ + --config configs/config_${RATE}.json \ + --gpu ${GPU} + + if [ $? -ne 0 ]; then + echo -e "${RED}训练失败: 缺失率 ${RATE}%${NC}" + exit 1 + fi + + echo "" +done + +# ============================================================================ +# 步骤3: 评估所有模型 +# ============================================================================ +echo -e "${YELLOW}步骤3: 评估所有模型${NC}" + +# 创建汇总结果文件 +SUMMARY_FILE="results/summary_results.txt" +mkdir -p results +echo "Missing_Rate,MAE,RMSE" > ${SUMMARY_FILE} + +for RATE in "${MISSING_RATES[@]}"; do + echo -e "${GREEN}========================================${NC}" + echo -e "${GREEN}评估缺失率 ${RATE}% 的模型${NC}" + echo -e "${GREEN}========================================${NC}" + + python evaluate.py \ + --config configs/config_${RATE}.json \ + --checkpoint results/traffic_${RATE}/best_model.pt \ + --gpu ${GPU} \ + --num_samples 10 + + if [ $? -ne 0 ]; then + echo -e "${RED}评估失败: 缺失率 ${RATE}%${NC}" + exit 1 + fi + + # 提取指标并添加到汇总文件 + METRICS_FILE="results/traffic_${RATE}/evaluation_metrics.txt" + if [ -f ${METRICS_FILE} ]; then + MAE=$(grep "MAE:" ${METRICS_FILE} | awk '{print $2}') + RMSE=$(grep "RMSE:" ${METRICS_FILE} | awk '{print $2}') + echo "${RATE},${MAE},${RMSE}" >> ${SUMMARY_FILE} + fi + + echo "" +done + +# ============================================================================ +# 步骤4: 生成汇总报告 +# ============================================================================ +echo -e "${YELLOW}步骤4: 生成汇总报告${NC}" + +# 创建汇总报告 +REPORT_FILE="results/final_report.txt" + +cat > ${REPORT_FILE} << EOF +======================================== +IEclaude 实验结果汇总 +======================================== + +实验配置: +- 数据集: LD2011_2014.txt +- 序列长度: 168 (1周) +- 模型: 隐式显式扩散模型 + * 隐式模块: TCN (扩张率 [1,2,4,8]) + * 显式模块: S4 (状态维度 64, 4层) +- 扩散步数: 200 +- 训练epochs: 100 + +实验结果: +======================================== +Missing Rate | MAE | RMSE +======================================== +EOF + +# 读取汇总文件并格式化输出 +tail -n +2 ${SUMMARY_FILE} | while IFS=',' read -r RATE MAE RMSE; do + printf "%-12s | %-6s | %-6s\n" "${RATE}%" "${MAE}" "${RMSE}" >> ${REPORT_FILE} +done + +echo "========================================" >> ${REPORT_FILE} +echo "" >> ${REPORT_FILE} +echo "目标指标 (参考):" >> ${REPORT_FILE} +echo "========================================" >> ${REPORT_FILE} +echo "Missing Rate | MAE | RMSE" >> ${REPORT_FILE} +echo "========================================" >> ${REPORT_FILE} +echo "20% | 0.272 | 0.389" >> ${REPORT_FILE} +echo "30% | 0.297 | 0.424" >> ${REPORT_FILE} +echo "40% | 0.334 | 0.477" >> ${REPORT_FILE} +echo "50% | 0.378 | 0.540" >> ${REPORT_FILE} +echo "60% | 0.450 | 0.655" >> ${REPORT_FILE} +echo "70% | 0.541 | 0.776" >> ${REPORT_FILE} +echo "80% | 0.732 | 1.049" >> ${REPORT_FILE} +echo "========================================" >> ${REPORT_FILE} + +# 显示汇总报告 +cat ${REPORT_FILE} + +# 保存汇总报告 +echo -e "\n${GREEN}汇总报告已保存: ${REPORT_FILE}${NC}" + +# ============================================================================ +# 完成 +# ============================================================================ +echo "" +echo -e "${GREEN}========================================${NC}" +echo -e "${GREEN}所有实验完成！${NC}" +echo -e "${GREEN}========================================${NC}" +echo -e "结果位置: ./results/" +echo -e "汇总报告: ${REPORT_FILE}" +echo "" diff --git a/IEclaude/train.py b/IEclaude/train.py new file mode 100644 index 0000000..e8e113c --- /dev/null +++ b/IEclaude/train.py @@ -0,0 +1,425 @@ +""" +IEclaude训练脚本 (Training Script) + +功能： + 训练隐式显式扩散模型用于交通数据插补 + +使用方法： + python train.py --config configs/config_20.json --gpu 0 + +超参数调优指南： + 见各个模块的注释和README文档 +""" + +import os +import sys +import json +import argparse +import torch +import torch.nn as nn +from torch.optim import Adam +from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR +from tqdm import tqdm +import numpy as np +import matplotlib.pyplot as plt +import shutil + +# 导入模块 +from models import IEDiffusionModel +from data import load_traffic_data, create_dataloader +from utils import DiffusionProcess, DiffusionLoss + + +def parse_args(): + """ + 解析命令行参数 + """ + parser = argparse.ArgumentParser(description='训练IEclaude模型') + + parser.add_argument('--config', type=str, required=True, + help='配置文件路径 (JSON格式)') + parser.add_argument('--gpu', type=int, default=0, + help='GPU编号 (0 for A40, 1 for A10)') + parser.add_argument('--resume', type=str, default=None, + help='从checkpoint恢复训练') + + return parser.parse_args() + + +def load_config(config_path): + """ + 加载配置文件 + + 参数： + config_path (str): 配置文件路径 + + 返回： + config (dict): 配置字典 + """ + with open(config_path, 'r', encoding='utf-8') as f: + config = json.load(f) + + return config + + +def set_seed(seed=42): + """ + 设置随机种子以保证可复现性 + + 参数： + seed (int): 随机种子 + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + + +def save_checkpoint(model, optimizer, epoch, loss, save_path, is_best=False): + """ + 保存checkpoint + + 参数： + model (nn.Module): 模型 + optimizer: 优化器 + epoch (int): 当前epoch + loss (float): 当前损失 + save_path (str): 保存路径 + is_best (bool): 是否是最佳模型 + """ + checkpoint = { + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'loss': loss + } + + torch.save(checkpoint, save_path) + + if is_best: + best_path = os.path.join(os.path.dirname(save_path), 'best_model.pt') + shutil.copyfile(save_path, best_path) + + +def train_epoch(model, dataloader, diffusion_process, loss_fn, optimizer, device, epoch, total_epochs): + """ + 训练一个epoch + + 参数： + model: 模型 + dataloader: 数据加载器 + diffusion_process: 扩散过程 + loss_fn: 损失函数 + optimizer: 优化器 + device: 设备 + epoch: 当前epoch + total_epochs: 总epoch数 + + 返回： + avg_loss (float): 平均损失 + """ + model.train() + total_loss = 0.0 + num_batches = len(dataloader) + + # 创建进度条 + pbar = tqdm(dataloader, desc=f'Epoch {epoch}/{total_epochs}') + + for batch_idx, batch in enumerate(pbar): + # 获取数据 + observed_data = batch['observed_data'].to(device) # [B, C, L] + mask = batch['mask'].to(device) # [B, C, L] + ground_truth = batch['ground_truth'].to(device) # [B, C, L] + + B, C, L = observed_data.shape + + # === 前向扩散：添加噪声 === + # 随机采样扩散步骤 + t = torch.randint(0, diffusion_process.T, (B,), device=device) + + # 添加噪声得到x_t + x_t, true_noise = diffusion_process.q_sample(ground_truth, t) + + # === 模型预测噪声 === + predicted_noise = model(x_t, observed_data, mask, t) + + # === 计算损失 === + loss = loss_fn(predicted_noise, true_noise, mask) + + # === 反向传播和优化 === + optimizer.zero_grad() + loss.backward() + + # 梯度裁剪：防止梯度爆炸 + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + # 累积损失 + total_loss += loss.item() + + # 更新进度条 + pbar.set_postfix({'loss': f'{loss.item():.6f}'}) + + # 计算平均损失 + avg_loss = total_loss / num_batches + + return avg_loss + + +def train(config, args): + """ + 主训练函数 + + 参数： + config (dict): 配置字典 + args: 命令行参数 + """ + print("=" * 80) + print("IEclaude 训练脚本") + print("=" * 80) + + # 设置随机种子 + set_seed(config.get('seed', 42)) + + # 设置设备 + device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') + print(f"\n使用设备: {device}") + + # 创建输出目录 + output_dir = config['train']['output_dir'] + os.makedirs(output_dir, exist_ok=True) + print(f"输出目录: {output_dir}") + + # 清除之前的训练结果 + if config['train'].get('clean_before_train', True): + print(f"\n清除之前的训练结果...") + for file in os.listdir(output_dir): + if file.endswith('.png') or file.endswith('.pt') or file.endswith('.txt'): + os.remove(os.path.join(output_dir, file)) + + # ======================================================================== + # 1. 加载数据 + # ======================================================================== + print("\n" + "=" * 80) + print("1. 加载数据") + print("=" * 80) + + train_data, test_data, scaler, data_info = load_traffic_data( + data_path=config['data']['data_path'], + seq_len=config['data']['seq_len'], + stride=config['data'].get('stride', config['data']['seq_len'] // 2), + train_ratio=config['data'].get('train_ratio', 0.7), + normalize_method=config['data'].get('normalize', 'standard'), + verbose=True + ) + + # 创建数据加载器 + train_loader = create_dataloader( + data=train_data, + missing_rate=config['train']['missing_rate'], + missing_pattern=config['train'].get('missing_pattern', 'random'), + batch_size=config['train']['batch_size'], + shuffle=True, + num_workers=config['train'].get('num_workers', 0) + ) + + print(f"\n训练集batch数量: {len(train_loader)}") + + # ======================================================================== + # 2. 创建模型 + # ======================================================================== + print("\n" + "=" * 80) + print("2. 创建模型") + print("=" * 80) + + model = IEDiffusionModel( + in_channels=data_info['num_channels'], + res_channels=config['model']['res_channels'], + skip_channels=config['model']['skip_channels'], + out_channels=data_info['num_channels'], + num_res_layers=config['model']['num_res_layers'], + dilation_cycle=config['model'].get('dilation_cycle', 10), + # 扩散步骤嵌入 + diffusion_step_embed_dim_in=config['diffusion']['embed_dim_in'], + diffusion_step_embed_dim_mid=config['diffusion']['embed_dim_mid'], + diffusion_step_embed_dim_out=config['diffusion']['embed_dim_out'], + # TCN参数 + tcn_channels=config['model']['tcn_channels'], + tcn_kernel_size=config['model']['tcn_kernel_size'], + tcn_dilation_rates=config['model']['tcn_dilation_rates'], + tcn_dropout=config['model'].get('tcn_dropout', 0.0), + # S4参数 + s4_d_state=config['model']['s4_d_state'], + s4_n_layers=config['model']['s4_n_layers'], + s4_l_max=config['data']['seq_len'], + s4_dropout=config['model'].get('s4_dropout', 0.0), + s4_bidirectional=config['model'].get('s4_bidirectional', True) + ).to(device) + + # 打印模型信息 + model_config = model.get_config() + print(f"\n模型配置:") + print(f" 输入通道数: {model_config['in_channels']}") + print(f" 残差通道数: {model_config['res_channels']}") + print(f" 残差层数: {model_config['num_res_layers']}") + print(f" TCN感受野: {model_config['tcn_config']['receptive_field']} 时间步") + print(f" S4状态维度: {model_config['s4_config']['d_state']}") + print(f" S4层数: {model_config['s4_config']['n_layers']}") + print(f" 参数总数: {model_config['num_parameters']:,}") + + # ======================================================================== + # 3. 创建扩散过程 + # ======================================================================== + print("\n" + "=" * 80) + print("3. 创建扩散过程") + print("=" * 80) + + diffusion_process = DiffusionProcess( + T=config['diffusion']['T'], + beta_0=config['diffusion']['beta_0'], + beta_T=config['diffusion']['beta_T'], + schedule=config['diffusion'].get('schedule', 'linear') + ) + + print(f" 扩散步数: {diffusion_process.T}") + print(f" 起始噪声: {diffusion_process.beta_0}") + print(f" 结束噪声: {diffusion_process.beta_T}") + print(f" 调度方式: {diffusion_process.schedule}") + + # 创建损失函数 + loss_fn = DiffusionLoss( + only_generate_missing=config['train'].get('only_generate_missing', True) + ) + + # ======================================================================== + # 4. 创建优化器和调度器 + # ======================================================================== + print("\n" + "=" * 80) + print("4. 创建优化器") + print("=" * 80) + + optimizer = Adam( + model.parameters(), + lr=config['train']['learning_rate'], + weight_decay=config['train'].get('weight_decay', 0.0) + ) + + # 学习率调度器 + scheduler_type = config['train'].get('scheduler', 'cosine') + if scheduler_type == 'cosine': + scheduler = CosineAnnealingLR( + optimizer, + T_max=config['train']['epochs'], + eta_min=config['train'].get('min_lr', 1e-6) + ) + elif scheduler_type == 'step': + scheduler = StepLR( + optimizer, + step_size=config['train'].get('step_size', 50), + gamma=config['train'].get('gamma', 0.5) + ) + else: + scheduler = None + + print(f" 优化器: Adam") + print(f" 初始学习率: {config['train']['learning_rate']}") + print(f" 调度器: {scheduler_type}") + + # ======================================================================== + # 5. 训练循环 + # ======================================================================== + print("\n" + "=" * 80) + print("5. 开始训练") + print("=" * 80) + + num_epochs = config['train']['epochs'] + save_interval = config['train'].get('save_interval', 10) + + train_losses = [] + best_loss = float('inf') + + for epoch in range(1, num_epochs + 1): + # 训练一个epoch + avg_loss = train_epoch( + model=model, + dataloader=train_loader, + diffusion_process=diffusion_process, + loss_fn=loss_fn, + optimizer=optimizer, + device=device, + epoch=epoch, + total_epochs=num_epochs + ) + + train_losses.append(avg_loss) + + # 更新学习率 + if scheduler is not None: + scheduler.step() + current_lr = optimizer.param_groups[0]['lr'] + else: + current_lr = config['train']['learning_rate'] + + # 打印信息 + print(f"\nEpoch {epoch}/{num_epochs} - Loss: {avg_loss:.6f} - LR: {current_lr:.6f}") + + # 保存checkpoint + if epoch % save_interval == 0: + checkpoint_path = os.path.join(output_dir, f'checkpoint_epoch_{epoch}.pt') + save_checkpoint(model, optimizer, epoch, avg_loss, checkpoint_path) + print(f" 已保存checkpoint: {checkpoint_path}") + + # 保存最佳模型 + if avg_loss < best_loss: + best_loss = avg_loss + best_path = os.path.join(output_dir, 'best_model.pt') + save_checkpoint(model, optimizer, epoch, avg_loss, best_path, is_best=True) + print(f" 已保存最佳模型: {best_path}") + + # ======================================================================== + # 6. 保存训练曲线 + # ======================================================================== + print("\n" + "=" * 80) + print("6. 保存训练曲线") + print("=" * 80) + + plt.figure(figsize=(10, 6)) + plt.plot(range(1, num_epochs + 1), train_losses, 'b-', label='Training Loss') + plt.xlabel('Epoch') + plt.ylabel('Loss') + plt.title('Training Loss Curve') + plt.legend() + plt.grid(True) + plt.savefig(os.path.join(output_dir, 'training_loss.png'), dpi=300, bbox_inches='tight') + plt.close() + + print(f" 训练曲线已保存: {os.path.join(output_dir, 'training_loss.png')}") + + # 保存损失值 + loss_file = os.path.join(output_dir, 'training_losses.txt') + with open(loss_file, 'w') as f: + f.write('Epoch,Loss\n') + for epoch, loss in enumerate(train_losses, 1): + f.write(f'{epoch},{loss:.6f}\n') + + print(f" 损失值已保存: {loss_file}") + + # ======================================================================== + # 7. 完成 + # ======================================================================== + print("\n" + "=" * 80) + print("训练完成！") + print("=" * 80) + print(f"最佳损失: {best_loss:.6f}") + print(f"模型和结果已保存到: {output_dir}") + + +if __name__ == '__main__': + # 解析命令行参数 + args = parse_args() + + # 加载配置 + config = load_config(args.config) + + # 开始训练 + train(config, args) diff --git a/IEclaude/utils/__init__.py b/IEclaude/utils/__init__.py new file mode 100644 index 0000000..534e19c --- /dev/null +++ b/IEclaude/utils/__init__.py @@ -0,0 +1,12 @@ +""" +IEclaude Utils Module + +包含扩散过程和其他实用工具 +""" + +from .diffusion import DiffusionProcess, DiffusionLoss + +__all__ = [ + 'DiffusionProcess', + 'DiffusionLoss' +] diff --git a/IEclaude/utils/diffusion.py b/IEclaude/utils/diffusion.py new file mode 100644 index 0000000..e023228 --- /dev/null +++ b/IEclaude/utils/diffusion.py @@ -0,0 +1,365 @@ +""" +扩散过程工具类 (Diffusion Process Utils) + +功能说明： + 实现DDPM (Denoising Diffusion Probabilistic Models) 的前向和反向扩散过程 + +核心概念： + 扩散模型通过逐步添加噪声将数据转换为纯噪声（前向过程）， + 然后训练神经网络学习逐步去噪（反向过程），从而实现数据生成和插补 + +前向扩散过程（加噪）： + q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I) + + 数学含义： + - x_0: 原始干净数据 + - x_t: 第t步的噪声数据 + - alpha_bar_t: 累积噪声系数 + - 随着t增大，x_t越来越接近纯高斯噪声 + +反向扩散过程（去噪）： + p_theta(x_{t-1} | x_t) = N(x_{t-1}; mu_theta(x_t, t), sigma_t^2 * I) + + 数学含义： + - 神经网络预测噪声 epsilon_theta(x_t, t) + - 使用预测的噪声计算 x_{t-1} + - 逐步从噪声恢复到干净数据 + +超参数调节指南： + 1. 扩散步数 (T): + - T越大，扩散过程越平滑，但计算量越大 + - 建议: 100-1000 + - 对于插补任务，200-500通常足够 + + 2. 噪声调度 (beta): + - beta_0: 起始噪声水平 [建议: 0.0001-0.001] + - beta_T: 结束噪声水平 [建议: 0.02-0.1] + - 线性调度: beta_t = beta_0 + (beta_T - beta_0) * t / T + - 余弦调度: 更平滑的噪声增长曲线 + + 3. 采样策略： + - DDPM: 标准采样，需要T步 + - DDIM: 加速采样，可以跳过一些步骤 + - 对于插补，建议使用DDPM以获得更好的质量 + +调优建议： + - 如果生成质量不好：增大T，调整beta范围 + - 如果训练不稳定：减小beta_T，增加warm-up + - 如果采样太慢：使用DDIM或减小T +""" + +import torch +import torch.nn as nn +import numpy as np + + +class DiffusionProcess: + """ + 扩散过程类 + + 功能： + - 定义前向扩散过程（加噪） + - 定义反向扩散过程（去噪） + - 计算扩散相关的参数 + + 参数： + T (int): 扩散步数 + beta_0 (float): 起始噪声水平 + beta_T (float): 结束噪声水平 + schedule (str): 噪声调度方式 ['linear', 'cosine'] + """ + def __init__(self, T=200, beta_0=0.0001, beta_T=0.02, schedule='linear'): + self.T = T + self.beta_0 = beta_0 + self.beta_T = beta_T + self.schedule = schedule + + # 计算噪声调度 + self.betas = self._get_noise_schedule() + + # 计算相关参数 + self.alphas = 1.0 - self.betas # alpha_t = 1 - beta_t + self.alpha_bars = torch.cumprod(self.alphas, dim=0) # alpha_bar_t = prod(alpha_i) for i=1..t + + # 计算其他有用的参数 + self.alpha_bars_prev = torch.cat([torch.tensor([1.0]), self.alpha_bars[:-1]]) # alpha_bar_{t-1} + self.sqrt_alpha_bars = torch.sqrt(self.alpha_bars) # sqrt(alpha_bar_t) + self.sqrt_one_minus_alpha_bars = torch.sqrt(1.0 - self.alpha_bars) # sqrt(1 - alpha_bar_t) + + # 反向过程的参数 + self.posterior_variance = ( + self.betas * (1.0 - self.alpha_bars_prev) / (1.0 - self.alpha_bars) + ) # 后验方差 + + def _get_noise_schedule(self): + """ + 计算噪声调度 + + 返回： + betas (torch.Tensor): 噪声调度 [T] + """ + if self.schedule == 'linear': + # 线性调度：beta从beta_0线性增长到beta_T + betas = torch.linspace(self.beta_0, self.beta_T, self.T) + + elif self.schedule == 'cosine': + # 余弦调度：更平滑的增长曲线 + # 参考: Improved Denoising Diffusion Probabilistic Models (Nichol & Dhariwal, 2021) + s = 0.008 # 小的偏移量 + steps = self.T + 1 + x = torch.linspace(0, self.T, steps) + alphas_bar = torch.cos(((x / self.T) + s) / (1 + s) * torch.pi * 0.5) ** 2 + alphas_bar = alphas_bar / alphas_bar[0] + betas = 1 - (alphas_bar[1:] / alphas_bar[:-1]) + betas = torch.clip(betas, 0.0001, 0.9999) # 裁剪到合理范围 + + else: + raise ValueError(f"未知的噪声调度方式: {self.schedule}") + + return betas + + def q_sample(self, x_0, t, noise=None): + """ + 前向扩散过程：从x_0采样x_t + + 数学公式： + x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon + 其中 epsilon ~ N(0, I) + + 参数： + x_0 (torch.Tensor): 原始数据 [B, C, L] + t (torch.Tensor): 时间步 [B] + noise (torch.Tensor, optional): 噪声 [B, C, L] + + 返回： + x_t (torch.Tensor): 加噪后的数据 [B, C, L] + noise (torch.Tensor): 添加的噪声 [B, C, L] + """ + # 如果未提供噪声，则采样标准高斯噪声 + if noise is None: + noise = torch.randn_like(x_0) + + # 获取对应时间步的参数 + sqrt_alpha_bar_t = self._extract(self.sqrt_alpha_bars, t, x_0.shape) + sqrt_one_minus_alpha_bar_t = self._extract(self.sqrt_one_minus_alpha_bars, t, x_0.shape) + + # 计算 x_t + x_t = sqrt_alpha_bar_t * x_0 + sqrt_one_minus_alpha_bar_t * noise + + return x_t, noise + + def p_sample(self, model, x_t, t, observed_data, mask): + """ + 反向扩散过程：从x_t采样x_{t-1} + + 数学公式： + 1. 使用模型预测噪声: epsilon_theta = model(x_t, observed_data, mask, t) + 2. 计算均值: mu = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1-alpha_bar_t)) * epsilon_theta) + 3. 采样: x_{t-1} = mu + sigma_t * z, 其中 z ~ N(0, I) + + 参数： + model (nn.Module): 去噪模型 + x_t (torch.Tensor): 当前噪声数据 [B, C, L] + t (torch.Tensor): 当前时间步 [B] + observed_data (torch.Tensor): 观测数据 [B, C, L] + mask (torch.Tensor): mask [B, C, L] + + 返回： + x_{t-1} (torch.Tensor): 去噪一步后的数据 [B, C, L] + """ + # 使用模型预测噪声 + predicted_noise = model(x_t, observed_data, mask, t) + + # 获取对应时间步的参数 + alpha_t = self._extract(self.alphas, t, x_t.shape) + alpha_bar_t = self._extract(self.alpha_bars, t, x_t.shape) + beta_t = self._extract(self.betas, t, x_t.shape) + + # 计算均值 + # mu_t = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * epsilon_theta) + coef1 = 1.0 / torch.sqrt(alpha_t) + coef2 = beta_t / torch.sqrt(1.0 - alpha_bar_t) + mean = coef1 * (x_t - coef2 * predicted_noise) + + # 计算方差 + posterior_var = self._extract(self.posterior_variance, t, x_t.shape) + + # 采样 x_{t-1} + if t[0] > 0: + # 如果不是最后一步，添加噪声 + noise = torch.randn_like(x_t) + x_t_minus_1 = mean + torch.sqrt(posterior_var) * noise + else: + # 如果是最后一步 (t=0)，不添加噪声 + x_t_minus_1 = mean + + return x_t_minus_1 + + def p_sample_loop(self, model, shape, observed_data, mask, device='cpu', verbose=False): + """ + 完整的反向扩散循环：从纯噪声逐步生成数据 + + 参数： + model (nn.Module): 去噪模型 + shape (tuple): 数据形状 (B, C, L) + observed_data (torch.Tensor): 观测数据 [B, C, L] + mask (torch.Tensor): mask [B, C, L] + device (str): 设备 + verbose (bool): 是否显示进度 + + 返回： + x_0 (torch.Tensor): 生成的数据 [B, C, L] + """ + B, C, L = shape + + # 从纯噪声开始 + x_t = torch.randn(B, C, L, device=device) + + # 逐步去噪 + for t_idx in reversed(range(self.T)): + # 当前时间步 + t = torch.full((B,), t_idx, device=device, dtype=torch.long) + + # 去噪一步 + x_t = self.p_sample(model, x_t, t, observed_data, mask) + + # 打印进度 + if verbose and (t_idx % 20 == 0 or t_idx == 0): + print(f" 去噪进度: {self.T - t_idx}/{self.T}") + + return x_t + + def _extract(self, a, t, x_shape): + """ + 从数组a中提取对应时间步t的值，并reshape到与x相同的形状 + + 参数： + a (torch.Tensor): 参数数组 [T] + t (torch.Tensor): 时间步 [B] + x_shape (tuple): 目标形状 + + 返回： + out (torch.Tensor): 提取并reshape后的值 + """ + batch_size = t.shape[0] + out = a.to(t.device)[t] # [B] + # Reshape到 [B, 1, 1, ...] 以便广播 + return out.view(batch_size, *([1] * (len(x_shape) - 1))) + + +class DiffusionLoss(nn.Module): + """ + 扩散模型损失函数 + + 功能： + 计算预测噪声和真实噪声之间的均方误差损失 + + 损失公式： + L = E_{t, x_0, epsilon} [ || epsilon - epsilon_theta(x_t, t) ||^2 ] + + 其中： + - epsilon: 真实添加的噪声 + - epsilon_theta: 模型预测的噪声 + - x_t: 通过前向扩散得到的噪声数据 + + 参数： + only_generate_missing (bool): 是否只对缺失位置计算损失 + - True: 只计算缺失位置的损失（适用于插补任务） + - False: 计算所有位置的损失（适用于生成任务） + """ + def __init__(self, only_generate_missing=True): + super().__init__() + self.only_generate_missing = only_generate_missing + + def forward(self, predicted_noise, true_noise, mask=None): + """ + 计算损失 + + 参数： + predicted_noise (torch.Tensor): 模型预测的噪声 [B, C, L] + true_noise (torch.Tensor): 真实噪声 [B, C, L] + mask (torch.Tensor, optional): mask [B, C, L] + 1=观测，0=缺失 + + 返回： + loss (torch.Tensor): 损失值（标量） + """ + if self.only_generate_missing and mask is not None: + # 只计算缺失位置的损失 + # mask中1表示观测，0表示缺失 + # 我们需要计算缺失位置，所以用 (1 - mask) + missing_mask = (1 - mask).float() + + # 计算缺失位置的MSE + loss = ((predicted_noise - true_noise) ** 2 * missing_mask).sum() + + # 归一化：除以缺失位置的数量 + num_missing = missing_mask.sum() + if num_missing > 0: + loss = loss / num_missing + else: + # 计算所有位置的MSE + loss = ((predicted_noise - true_noise) ** 2).mean() + + return loss + + +# ============================================================================ +# 使用示例和调试代码 +# ============================================================================ + +if __name__ == "__main__": + """ + 测试扩散过程 + """ + print("=" * 80) + print("扩散过程测试") + print("=" * 80) + + # 创建扩散过程 + diffusion = DiffusionProcess(T=200, beta_0=0.0001, beta_T=0.02, schedule='linear') + + print(f"\n扩散过程配置:") + print(f" 扩散步数 T: {diffusion.T}") + print(f" 起始噪声 beta_0: {diffusion.beta_0}") + print(f" 结束噪声 beta_T: {diffusion.beta_T}") + print(f" 调度方式: {diffusion.schedule}") + + # 创建模拟数据 + batch_size = 4 + channels = 370 + seq_len = 168 + + x_0 = torch.randn(batch_size, channels, seq_len) + print(f"\n原始数据 x_0: {x_0.shape}") + + # 测试前向扩散 + print(f"\n测试前向扩散...") + t = torch.randint(0, diffusion.T, (batch_size,)) + x_t, noise = diffusion.q_sample(x_0, t) + print(f" 时间步 t: {t}") + print(f" 加噪数据 x_t: {x_t.shape}") + print(f" 噪声: {noise.shape}") + + # 可视化不同时间步的噪声水平 + print(f"\n不同时间步的噪声系数:") + test_steps = [0, 50, 100, 150, 199] + for step in test_steps: + sqrt_alpha_bar = diffusion.sqrt_alpha_bars[step] + sqrt_one_minus_alpha_bar = diffusion.sqrt_one_minus_alpha_bars[step] + print(f" t={step:3d}: sqrt(alpha_bar)={sqrt_alpha_bar:.4f}, sqrt(1-alpha_bar)={sqrt_one_minus_alpha_bar:.4f}") + + # 测试损失函数 + print(f"\n测试损失函数...") + loss_fn = DiffusionLoss(only_generate_missing=True) + + # 创建mask + mask = torch.rand(batch_size, channels, seq_len) > 0.3 + predicted_noise = torch.randn_like(noise) + + loss = loss_fn(predicted_noise, noise, mask) + print(f" 损失值: {loss.item():.6f}") + + print("\n" + "=" * 80) + print("测试完成！") + print("=" * 80) From e762e0807ec7ae2d59dcbc38f34a91a7e9f1f45a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Nov 2025 14:54:42 +0000 Subject: [PATCH 3/3] Fix S4 bidirectional mode dimension mismatch and add GPU checker - Fix: S4Layer skip connection must be added before bidirectional concatenation - Add: check_gpu.py script to detect available GPUs and provide usage suggestions - Fix resolves RuntimeError: The size of tensor a (512) must match the size of tensor b (256) --- IEclaude/check_gpu.py | 69 ++++++++++++++++++++++++++++++++++ IEclaude/models/s4_explicit.py | 11 ++++-- 2 files changed, 76 insertions(+), 4 deletions(-) create mode 100644 IEclaude/check_gpu.py diff --git a/IEclaude/check_gpu.py b/IEclaude/check_gpu.py new file mode 100644 index 0000000..66113fe --- /dev/null +++ b/IEclaude/check_gpu.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +""" +GPU检测脚本 + +检查系统中可用的GPU设备 +""" + +import torch + +print("=" * 80) +print("GPU 设备检测") +print("=" * 80) + +# 检查CUDA是否可用 +if torch.cuda.is_available(): + print(f"\n✓ CUDA可用") + print(f" CUDA版本: {torch.version.cuda}") + + # 获取GPU数量 + num_gpus = torch.cuda.device_count() + print(f" 可用GPU数量: {num_gpus}") + + # 列出所有GPU + print(f"\nGPU详细信息:") + for i in range(num_gpus): + print(f"\n GPU {i}:") + print(f" 名称: {torch.cuda.get_device_name(i)}") + + # 获取显存信息 + props = torch.cuda.get_device_properties(i) + total_memory = props.total_memory / 1024**3 # 转换为GB + print(f" 显存: {total_memory:.2f} GB") + + # 当前显存使用情况 + if i < num_gpus: + torch.cuda.set_device(i) + allocated = torch.cuda.memory_allocated(i) / 1024**3 + reserved = torch.cuda.memory_reserved(i) / 1024**3 + print(f" 已分配: {allocated:.2f} GB") + print(f" 已保留: {reserved:.2f} GB") + + # 使用建议 + print(f"\n" + "=" * 80) + print("使用建议:") + print("=" * 80) + + for i in range(num_gpus): + gpu_name = torch.cuda.get_device_name(i) + props = torch.cuda.get_device_properties(i) + total_memory = props.total_memory / 1024**3 + + print(f"\nGPU {i} ({gpu_name}):") + print(f" 训练命令: python train.py --config configs/config_20.json --gpu {i}") + print(f" 或使用脚本: ./run_all.sh --gpu {i}") + + # 根据显存给出batch_size建议 + if total_memory < 12: + print(f" 建议batch_size: 4 (显存较小)") + elif total_memory < 24: + print(f" 建议batch_size: 8 (默认)") + else: + print(f" 建议batch_size: 16 (显存充足)") + +else: + print(f"\n✗ CUDA不可用") + print(f" 将使用CPU进行训练（速度会很慢）") + print(f" 训练命令: python train.py --config configs/config_20.json") + +print("\n" + "=" * 80) diff --git a/IEclaude/models/s4_explicit.py b/IEclaude/models/s4_explicit.py index 3c1211c..7797a62 100644 --- a/IEclaude/models/s4_explicit.py +++ b/IEclaude/models/s4_explicit.py @@ -371,6 +371,10 @@ def forward(self, u): y_f = u_f * k_f.unsqueeze(0) # [B, d_model, L_fft] y = torch.fft.irfft(y_f, n=2*L)[..., :L] # [B, d_model, L] + # === 添加skip connection (D矩阵) - 必须在双向拼接前 === + D = self.kernel.D.unsqueeze(0).unsqueeze(-1) # [1, d_model, 1] + y = y + u * D # [B, d_model, L] + # === 双向处理 === if self.bidirectional: # 后向卷积：反转序列 @@ -382,13 +386,12 @@ def forward(self, u): y_b = torch.fft.irfft(y_b_f, n=2*L)[..., :L] y_b = torch.flip(y_b, dims=[-1]) # 反转回来 + # 后向也添加skip connection + y_b = y_b + u * D # [B, d_model, L] + # 拼接前向和后向 y = torch.cat([y, y_b], dim=1) # [B, 2*d_model, L] - # === 添加skip connection (D矩阵) === - D = self.kernel.D.unsqueeze(0).unsqueeze(-1) # [1, d_model, 1] - y = y + u * D # [B, d_model, L] 或 [B, 2*d_model, L] - # === 激活和Dropout === y = self.activation(y) y = self.dropout(y)