diff --git a/COMMANDS.md b/COMMANDS.md index 27920dc..c90fd6a 100644 --- a/COMMANDS.md +++ b/COMMANDS.md @@ -2,10 +2,22 @@ ## 训练 +### WINDOWS+GPU + ```bash python src/train.py experiment=PPG_FieldStudy_Base callbacks=PPG_FieldStudy logger=tensorboard + python src/train.py experiment=PPG_FieldStudy_Base callbacks=PPG_FieldStudy logger=tensorboard trainer=gpu -python src/train.py experiment=PPG_FieldStudy_Base logger=tensorboard callbacks=PPG_FieldStudy trainer=gpu debug=fdr + +python src/train.py experiment=PPG_FieldStudy_Base logger=tensorboard callbacks=PPG_FieldStudy trainer=gpu debug=fdr +``` + +### WINDOWS+CPU + +```powershell +python src/train.py experiment=PPG_FieldStudy_Base callbacks=PPG_FieldStudy logger=tensorboard trainer.accelerator=cpu + +python src/train.py experiment=PPG_FieldStudy_LSTM callbacks=PPG_FieldStudy logger=tensorboard trainer.accelerator=cpu ``` ## 实验结果查看 @@ -21,5 +33,3 @@ tensorboard --logdir=./logs/train/runs/ ```bash python src/eval.py ckpt_path="D:\lightning-hydra-template\logs\train\runs\2024-10-06_11-35-20\checkpoints\epoch_005.ckpt" ``` - - diff --git a/README.md b/README.md index 6146fa4..369496a 100644 --- a/README.md +++ b/README.md @@ -9,3 +9,9 @@ Related models: 3. Transformer 4. Mamba 5. KAN + +## Datasets + +### PPG_FieldStudy + +https://www.kaggle.com/datasets/dishantvishwakarma/ppg-dataset-shared/data diff --git a/configs/experiment/PPG_FieldStudy_GRU.yaml b/configs/experiment/PPG_FieldStudy_GRU.yaml new file mode 100644 index 0000000..16ed7cf --- /dev/null +++ b/configs/experiment/PPG_FieldStudy_GRU.yaml @@ -0,0 +1,46 @@ +# @package _global_ + +# to execute this experiment run: +# python train.py experiment=example + +defaults: + - override /data: PPG_FieldStudy + - override /model: PPGGRU + - override /callbacks: default + - override /trainer: default + +# all parameters below will be merged with parameters from default configurations set above +# this allows you to overwrite only specified parameters + +tags: ["PPG_FieldStudy", "PPGGRU"] + +seed: 12345 + +trainer: + min_epochs: 10 + max_epochs: 50 + gradient_clip_val: 0.5 + accelerator: "gpu" + devices: 1 + callbacks: [progress_bar] + enable_checkpointing: True + +model: + input_dim: 14 + hidden_dim: 128 + num_layers: 2 + bidirectional: True + dropout: 0.3 + lr: 1e-3 + scheduler_step_size: 10 + scheduler_gamma: 0.1 + +data: + batch_size: 64 + +logger: + wandb: + tags: ${tags} + group: "PPG_FieldStudy" + aim: + experiment: "PPG_FieldStudy" diff --git a/configs/experiment/PPG_FieldStudy_LSTM.yaml b/configs/experiment/PPG_FieldStudy_LSTM.yaml new file mode 100644 index 0000000..4d013c3 --- /dev/null +++ b/configs/experiment/PPG_FieldStudy_LSTM.yaml @@ -0,0 +1,46 @@ +# @package _global_ + +# to execute this experiment run: +# python train.py experiment=example + +defaults: + - override /data: PPG_FieldStudy + - override /model: PPGLSTM + - override /callbacks: default + - override /trainer: default + +# all parameters below will be merged with parameters from default configurations set above +# this allows you to overwrite only specified parameters + +tags: ["PPG_FieldStudy", "PPGLSTM"] + +seed: 12345 + +trainer: + min_epochs: 10 + max_epochs: 50 + gradient_clip_val: 0.5 + accelerator: "gpu" + devices: 1 + callbacks: [progress_bar] + enable_checkpointing: True + +model: + input_dim: 14 + hidden_dim: 128 + num_layers: 2 + bidirectional: True + dropout: 0.3 + lr: 1e-3 + scheduler_step_size: 10 + scheduler_gamma: 0.1 + +data: + batch_size: 64 + +logger: + wandb: + tags: ${tags} + group: "PPG_FieldStudy" + aim: + experiment: "PPG_FieldStudy" diff --git a/configs/experiment/PPG_FieldStudy_Mamba.yaml b/configs/experiment/PPG_FieldStudy_Mamba.yaml new file mode 100644 index 0000000..4aa2583 --- /dev/null +++ b/configs/experiment/PPG_FieldStudy_Mamba.yaml @@ -0,0 +1,47 @@ +# @package _global_ + +# to execute this experiment run: +# python train.py experiment=example + +defaults: + - override /data: PPG_FieldStudy + - override /model: PPGMamba + - override /callbacks: default + - override /trainer: default + +# all parameters below will be merged with parameters from default configurations set above +# this allows you to overwrite only specified parameters + +tags: ["PPG_FieldStudy", "PPGMamba"] + +seed: 12345 + +trainer: + min_epochs: 10 + max_epochs: 50 + gradient_clip_val: 0.5 + accelerator: "gpu" + devices: 1 + callbacks: [progress_bar] + enable_checkpointing: True + +model: + input_dim: 14 + d_model: 128 + d_state: 64 + d_conv: 4 + expand: 2 + dropout: 0.3 + lr: 1e-3 + scheduler_step_size: 10 + scheduler_gamma: 0.1 + +data: + batch_size: 64 + +logger: + wandb: + tags: ${tags} + group: "PPG_FieldStudy" + aim: + experiment: "PPG_FieldStudy" diff --git a/configs/experiment/PPG_FieldStudy_Transformer.yaml b/configs/experiment/PPG_FieldStudy_Transformer.yaml new file mode 100644 index 0000000..aec7a06 --- /dev/null +++ b/configs/experiment/PPG_FieldStudy_Transformer.yaml @@ -0,0 +1,47 @@ +# @package _global_ + +# to execute this experiment run: +# python train.py experiment=example + +defaults: + - override /data: PPG_FieldStudy + - override /model: PPGTransformer + - override /callbacks: default + - override /trainer: default + +# all parameters below will be merged with parameters from default configurations set above +# this allows you to overwrite only specified parameters + +tags: ["PPG_FieldStudy", "PPGTransformer"] + +seed: 12345 + +trainer: + min_epochs: 10 + max_epochs: 50 + gradient_clip_val: 0.5 + accelerator: "gpu" + devices: 1 + callbacks: [progress_bar] + enable_checkpointing: True + +model: + input_dim: 14 + d_model: 128 + nhead: 8 + num_encoder_layers: 4 + dim_feedforward: 256 + dropout: 0.1 + lr: 1e-3 + scheduler_step_size: 10 + scheduler_gamma: 0.1 + +data: + batch_size: 64 + +logger: + wandb: + tags: ${tags} + group: "PPG_FieldStudy" + aim: + experiment: "PPG_FieldStudy" diff --git a/configs/model/PPGGRU.yaml b/configs/model/PPGGRU.yaml new file mode 100644 index 0000000..62143e0 --- /dev/null +++ b/configs/model/PPGGRU.yaml @@ -0,0 +1,10 @@ +_target_: src.models.PPGGRUModule.PPGGRUModule + +input_dim: 14 +hidden_dim: 128 +num_layers: 2 +bidirectional: True +dropout: 0.3 +lr: 1e-3 +scheduler_step_size: 10 +scheduler_gamma: 0.1 diff --git a/configs/model/PPGLSTM.yaml b/configs/model/PPGLSTM.yaml new file mode 100644 index 0000000..5e1b7e0 --- /dev/null +++ b/configs/model/PPGLSTM.yaml @@ -0,0 +1,10 @@ +_target_: src.models.PPGLSTMModule.PPGLSTMModule + +input_dim: 14 +hidden_dim: 128 +num_layers: 2 +bidirectional: True +dropout: 0.3 +lr: 1e-3 +scheduler_step_size: 10 +scheduler_gamma: 0.1 diff --git a/configs/model/PPGMamba.yaml b/configs/model/PPGMamba.yaml new file mode 100644 index 0000000..d5a930a --- /dev/null +++ b/configs/model/PPGMamba.yaml @@ -0,0 +1,11 @@ +_target_: src.models.PPGMambaModule.PPGMambaModule + +input_dim: 14 +d_model: 128 +d_state: 64 +d_conv: 4 +expand: 2 +dropout: 0.3 +lr: 1e-3 +scheduler_step_size: 10 +scheduler_gamma: 0.1 diff --git a/configs/model/PPGTransformer.yaml b/configs/model/PPGTransformer.yaml new file mode 100644 index 0000000..55ecd8c --- /dev/null +++ b/configs/model/PPGTransformer.yaml @@ -0,0 +1,11 @@ +_target_: src.models.PPGTransformerModule.PPGTransformerModule + +input_dim: 14 +d_model: 128 +nhead: 8 +num_encoder_layers: 4 +dim_feedforward: 256 +dropout: 0.1 +lr: 1e-3 +scheduler_step_size: 10 +scheduler_gamma: 0.1 diff --git a/src/data/PPGFieldStudyDatamodule.py b/src/data/PPGFieldStudyDatamodule.py index 6f01266..aefc87f 100644 --- a/src/data/PPGFieldStudyDatamodule.py +++ b/src/data/PPGFieldStudyDatamodule.py @@ -44,8 +44,8 @@ def __init__( # Validate labels length assert ( - len(labels) == self.num_windows - ), f"Number of labels ({len(labels)}) does not match number of windows ({self.num_windows})" + len(labels) >= self.num_windows + ), f"Number of labels ({len(labels)}) is less than number of windows ({self.num_windows})." # Normalize signals using StandardScaler self.scalers = {} @@ -61,7 +61,7 @@ def __init__( self.scalers[location][sensor] = scaler self.data = data - self.labels = labels + self.labels = labels[:self.num_windows] # Ensure labels match the number of windows def __len__(self): return self.num_windows @@ -93,21 +93,21 @@ def __getitem__(self, idx): # Concatenate all sensor data along the feature dimension chest_features = np.concatenate( [features["chest"][sensor] for sensor in features["chest"]], axis=1 - ) # Shape: [256, 8] + ) # Shape: [window_size, 8] wrist_features = np.concatenate( [features["wrist"][sensor] for sensor in features["wrist"]], axis=1 - ) # Shape: [256, 6] + ) # Shape: [window_size, 6] combined_features = np.concatenate( [chest_features, wrist_features], axis=1 - ) # Shape: [256, 14] + ) # Shape: [window_size, 14] if self.transform: combined_features = self.transform(combined_features) - # Aggregate window into single vector by computing the mean across the window_size dimension - aggregated_features = combined_features.mean(axis=0) # Shape: [14] + # 保留整个窗口的特征，不进行汇聚 + aggregated_features = combined_features # Shape: [window_size, 14] - # Ensure labels are aligned with windows + # 获取对应的标签 label = self.labels[idx] return torch.tensor(aggregated_features, dtype=torch.float32), torch.tensor( diff --git a/src/models/PPGBaseModule.py b/src/models/PPGBaseModule.py index 0e8148d..e691389 100644 --- a/src/models/PPGBaseModule.py +++ b/src/models/PPGBaseModule.py @@ -4,14 +4,14 @@ from lightning import LightningModule from torch.optim import Adam from torch.optim.lr_scheduler import StepLR -from torchmetrics import MeanSquaredError from torchmetrics import MeanMetric class PPGBaseModule(LightningModule): def __init__( self, - input_dim: int, + input_dim: int = 14, # 每个时间步的特征数 + window_size: int = 256, # 时间步数 hidden_dim: int = 128, lr: float = 1e-3, scheduler_step_size: int = 10, @@ -20,7 +20,8 @@ def __init__( """ LightningModule for PPG-based physiological indicator prediction. - :param input_dim: Number of input features. + :param input_dim: Number of input features per time step. + :param window_size: Number of time steps per window. :param hidden_dim: Number of hidden units. :param lr: Learning rate. :param scheduler_step_size: Step size for learning rate scheduler. @@ -29,22 +30,19 @@ def __init__( super().__init__() self.save_hyperparameters() - # Define the network architecture - self.encoder = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(hidden_dim, hidden_dim // 2), - nn.ReLU(), - nn.Dropout(0.3), - ) + # 定义1D卷积网络 + self.conv1 = nn.Conv1d(in_channels=input_dim, out_channels=hidden_dim, kernel_size=3, padding=1) + self.bn1 = nn.BatchNorm1d(hidden_dim) + self.conv2 = nn.Conv1d(in_channels=hidden_dim, out_channels=hidden_dim*2, kernel_size=3, padding=1) + self.bn2 = nn.BatchNorm1d(hidden_dim*2) + self.pool = nn.MaxPool1d(kernel_size=2) - self.regressor = nn.Sequential( - nn.Linear(hidden_dim // 2, hidden_dim // 4), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(hidden_dim // 4, 1), # Single output - ) + # 计算池化后的时间步数 + pooled_size = window_size // 2 # 因为进行了一个池化层，时间步数减半 + + self.fc1 = nn.Linear((hidden_dim*2) * pooled_size, hidden_dim) + self.dropout = nn.Dropout(0.5) + self.fc2 = nn.Linear(hidden_dim, 1) # Loss function self.criterion = nn.MSELoss() @@ -58,12 +56,18 @@ def forward(self, x): """ Forward pass. - :param x: Input tensor of shape [batch_size, 14] + :param x: Input tensor of shape [batch_size, window_size, 14] :return: Predicted tensor of shape [batch_size] """ - encoded = self.encoder(x) # [batch_size, hidden_dim//2] - output = self.regressor(encoded) # [batch_size, 1] - return output.squeeze(1) # [batch_size] + x = x.permute(0, 2, 1) # 转换为 [batch_size, 14, window_size] 以适应 Conv1d + x = F.relu(self.bn1(self.conv1(x))) # [batch_size, hidden_dim, window_size] + x = F.relu(self.bn2(self.conv2(x))) # [batch_size, hidden_dim*2, window_size] + x = self.pool(x) # [batch_size, hidden_dim*2, window_size//2] + x = x.view(x.size(0), -1) # 展平成 [batch_size, hidden_dim*2 * (window_size//2)] + x = F.relu(self.fc1(x)) # [batch_size, hidden_dim] + x = self.dropout(x) + x = self.fc2(x).squeeze(1) # [batch_size] + return x def training_step(self, batch, batch_idx): """ @@ -73,7 +77,7 @@ def training_step(self, batch, batch_idx): :param batch_idx: Batch index. :return: Loss value. """ - x, y = batch # x: [batch_size, 14], y: [batch_size] + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] y_pred = self.forward(x) # [batch_size] loss = self.criterion(y_pred, y) @@ -93,7 +97,7 @@ def validation_step(self, batch, batch_idx): :param batch: Batch of data. :param batch_idx: Batch index. """ - x, y = batch # x: [batch_size, 14], y: [batch_size] + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] y_pred = self.forward(x) # [batch_size] loss = self.criterion(y_pred, y) @@ -109,7 +113,7 @@ def test_step(self, batch, batch_idx): :param batch: Batch of data. :param batch_idx: Batch index. """ - x, y = batch # x: [batch_size, 14], y: [batch_size] + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] y_pred = self.forward(x) # [batch_size] loss = self.criterion(y_pred, y) diff --git a/src/models/PPGGRUModule.py b/src/models/PPGGRUModule.py index e69de29..c34ebe1 100644 --- a/src/models/PPGGRUModule.py +++ b/src/models/PPGGRUModule.py @@ -0,0 +1,164 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from lightning import LightningModule +from torch.optim import Adam +from torch.optim.lr_scheduler import StepLR +from torchmetrics import MeanMetric + +class PPGGRUModule(LightningModule): + def __init__( + self, + input_dim: int = 14, # 每个时间步的特征数 + hidden_dim: int = 128, # GRU隐藏状态的维度 + num_layers: int = 2, # GRU层数 + bidirectional: bool = True, # 是否使用双向GRU + dropout: float = 0.3, # GRU和全连接层之间的Dropout比例 + lr: float = 1e-3, # 学习率 + scheduler_step_size: int = 10,# 学习率调度器的步长 + scheduler_gamma: float = 0.1, # 学习率调度器的gamma值 + ): + """ + LightningModule 使用 GRU 进行 PPG 数据预测。 + + :param input_dim: 每个时间步的输入特征数。 + :param hidden_dim: GRU隐藏状态的维度。 + :param num_layers: GRU的层数。 + :param bidirectional: 是否使用双向GRU。 + :param dropout: Dropout比例。 + :param lr: 学习率。 + :param scheduler_step_size: 学习率调度器的步长。 + :param scheduler_gamma: 学习率调度器的gamma值。 + """ + super().__init__() + self.save_hyperparameters() + + # 输入投影层：将输入特征从 input_dim 映射到 hidden_dim + self.input_projection = nn.Linear(input_dim, hidden_dim) + + # GRU层 + self.gru = nn.GRU( + input_size=hidden_dim, + hidden_size=hidden_dim, + num_layers=num_layers, + batch_first=True, + bidirectional=bidirectional, + dropout=dropout if num_layers > 1 else 0.0 # 只有多层时才使用Dropout + ) + + # 计算GRU输出的特征维度 + gru_output_dim = hidden_dim * 2 if bidirectional else hidden_dim + + # 回归头：全连接层 + self.fc1 = nn.Linear(gru_output_dim, gru_output_dim // 2) + self.dropout = nn.Dropout(dropout) + self.fc2 = nn.Linear(gru_output_dim // 2, 1) # 单输出 + + # 损失函数 + self.criterion = nn.MSELoss() + + # 评价指标 + self.train_mse = MeanMetric() + self.val_mse = MeanMetric() + self.test_mse = MeanMetric() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + 前向传播。 + + :param x: 输入张量，形状为 [batch_size, seq_len, input_dim] + :return: 预测张量，形状为 [batch_size] + """ + # 输入投影 + x = self.input_projection(x) # [batch_size, seq_len, hidden_dim] + + # GRU层 + # h_0的形状为 [num_layers * num_directions, batch_size, hidden_dim] + h_0 = torch.zeros( + self.hparams.num_layers * (2 if self.hparams.bidirectional else 1), + x.size(0), + self.hparams.hidden_dim, + device=x.device + ) + gru_out, h_n = self.gru(x, h_0) # gru_out: [batch_size, seq_len, hidden_dim * num_directions] + + # 处理双向GRU的隐藏状态 + if self.hparams.bidirectional: + # 取最后一层的正向和反向隐藏状态 + h_forward = h_n[-2, :, :] # [batch_size, hidden_dim] + h_backward = h_n[-1, :, :] # [batch_size, hidden_dim] + h = torch.cat((h_forward, h_backward), dim=1) # [batch_size, hidden_dim * 2] + else: + h = h_n[-1, :, :] # [batch_size, hidden_dim] + + # 全连接层 + x = F.relu(self.fc1(h)) # [batch_size, hidden_dim] + x = self.dropout(x) + x = self.fc2(x).squeeze(1) # [batch_size] + + return x + + def training_step(self, batch, batch_idx): + """ + 训练步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + :return: 损失值。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.train_mse(y_pred, y) + self.log("train/mse", self.train_mse, on_step=False, on_epoch=True, prog_bar=True) + + return loss + + def validation_step(self, batch, batch_idx): + """ + 验证步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.val_mse(y_pred, y) + self.log("val/mse", self.val_mse, on_step=False, on_epoch=True, prog_bar=True) + + def test_step(self, batch, batch_idx): + """ + 测试步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("test/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.test_mse(y_pred, y) + self.log("test/mse", self.test_mse, on_step=False, on_epoch=True, prog_bar=True) + + def configure_optimizers(self): + """ + 配置优化器和学习率调度器。 + + :return: 优化器和调度器配置。 + """ + optimizer = Adam(self.parameters(), lr=self.hparams.lr) + scheduler = StepLR( + optimizer, + step_size=self.hparams.scheduler_step_size, + gamma=self.hparams.scheduler_gamma, + ) + return [optimizer], [scheduler] diff --git a/src/models/PPGLSTMModule.py b/src/models/PPGLSTMModule.py index e69de29..47a0e43 100644 --- a/src/models/PPGLSTMModule.py +++ b/src/models/PPGLSTMModule.py @@ -0,0 +1,152 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from lightning import LightningModule +from torch.optim import Adam +from torch.optim.lr_scheduler import StepLR +from torchmetrics import MeanMetric + +class PPGLSTMModule(LightningModule): + def __init__( + self, + input_dim: int = 14, # 每个时间步的特征数 + hidden_dim: int = 128, # LSTM隐藏状态的维度 + num_layers: int = 2, # LSTM层数 + bidirectional: bool = True, # 是否使用双向LSTM + dropout: float = 0.3, # LSTM和全连接层之间的Dropout比例 + lr: float = 1e-3, # 学习率 + scheduler_step_size: int = 10, # 学习率调度器的步长 + scheduler_gamma: float = 0.1, # 学习率调度器的gamma值 + ): + """ + LightningModule 使用 LSTM 进行 PPG 数据预测。 + + :param input_dim: 每个时间步的输入特征数。 + :param hidden_dim: LSTM隐藏状态的维度。 + :param num_layers: LSTM的层数。 + :param bidirectional: 是否使用双向LSTM。 + :param dropout: Dropout的比例。 + :param lr: 学习率。 + :param scheduler_step_size: 学习率调度器的步长。 + :param scheduler_gamma: 学习率调度器的gamma值。 + """ + super().__init__() + self.save_hyperparameters() + + # 定义LSTM层 + self.lstm = nn.LSTM( + input_size=input_dim, + hidden_size=hidden_dim, + num_layers=num_layers, + batch_first=True, + bidirectional=bidirectional, + dropout=dropout if num_layers > 1 else 0.0 # 只有多层时才使用Dropout + ) + + # 计算LSTM输出的特征维度 + lstm_output_dim = hidden_dim * 2 if bidirectional else hidden_dim + + # 定义全连接层 + self.fc1 = nn.Linear(lstm_output_dim, hidden_dim) + self.dropout = nn.Dropout(dropout) + self.fc2 = nn.Linear(hidden_dim, 1) # 单输出 + + # 损失函数 + self.criterion = nn.MSELoss() + + # 评价指标 + self.train_mse = MeanMetric() + self.val_mse = MeanMetric() + self.test_mse = MeanMetric() + + def forward(self, x): + """ + 前向传播。 + + :param x: 输入张量，形状为 [batch_size, window_size, 14] + :return: 预测张量，形状为 [batch_size] + """ + # LSTM的输入形状为 [batch_size, window_size, input_dim] + lstm_out, (h_n, c_n) = self.lstm(x) + # h_n的形状为 [num_layers * num_directions, batch_size, hidden_dim] + + # 如果使用双向LSTM，连接最后一层的正向和反向隐藏状态 + if self.hparams.bidirectional: + # 取最后一层的正向和反向隐藏状态 + h_forward = h_n[-2, :, :] # [batch_size, hidden_dim] + h_backward = h_n[-1, :, :] # [batch_size, hidden_dim] + h = torch.cat((h_forward, h_backward), dim=1) # [batch_size, hidden_dim * 2] + else: + h = h_n[-1, :, :] # [batch_size, hidden_dim] + + # 通过全连接层 + x = F.relu(self.fc1(h)) # [batch_size, hidden_dim] + x = self.dropout(x) + x = self.fc2(x).squeeze(1) # [batch_size] + + return x + + def training_step(self, batch, batch_idx): + """ + 训练步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + :return: 损失值。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.train_mse(y_pred, y) + self.log("train/mse", self.train_mse, on_step=False, on_epoch=True, prog_bar=True) + + return loss + + def validation_step(self, batch, batch_idx): + """ + 验证步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.val_mse(y_pred, y) + self.log("val/mse", self.val_mse, on_step=False, on_epoch=True, prog_bar=True) + + def test_step(self, batch, batch_idx): + """ + 测试步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("test/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.test_mse(y_pred, y) + self.log("test/mse", self.test_mse, on_step=False, on_epoch=True, prog_bar=True) + + def configure_optimizers(self): + """ + 配置优化器和学习率调度器。 + + :return: 优化器和调度器配置。 + """ + optimizer = Adam(self.parameters(), lr=self.hparams.lr) + scheduler = StepLR( + optimizer, + step_size=self.hparams.scheduler_step_size, + gamma=self.hparams.scheduler_gamma, + ) + return [optimizer], [scheduler] diff --git a/src/models/PPGMambaModule.py b/src/models/PPGMambaModule.py index e69de29..32c1ca0 100644 --- a/src/models/PPGMambaModule.py +++ b/src/models/PPGMambaModule.py @@ -0,0 +1,156 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from lightning import LightningModule +from torch.optim import Adam +from torch.optim.lr_scheduler import StepLR +from torchmetrics import MeanMetric +from mamba_ssm import Mamba2 # 确保已安装 mamba_ssm 并正确导入 + +class PPGMambaModule(LightningModule): + def __init__( + self, + input_dim: int = 14, # 每个时间步的特征数 + d_model: int = 128, # Mamba2 模型的特征维度 + d_state: int = 64, # SSM 状态扩展因子 + d_conv: int = 4, # 局部卷积宽度 + expand: int = 2, # 块扩展因子 + dropout: float = 0.3, # Dropout 比例 + lr: float = 1e-3, # 学习率 + scheduler_step_size: int = 10,# 学习率调度器的步长 + scheduler_gamma: float = 0.1, # 学习率调度器的 gamma 值 + ): + """ + LightningModule 使用 Mamba2 进行 PPG 数据预测。 + + :param input_dim: 每个时间步的输入特征数。 + :param d_model: Mamba2 模型的特征维度。 + :param d_state: SSM 状态扩展因子，通常为 64 或 128。 + :param d_conv: 局部卷积宽度。 + :param expand: 块扩展因子。 + :param dropout: Dropout 比例。 + :param lr: 学习率。 + :param scheduler_step_size: 学习率调度器的步长。 + :param scheduler_gamma: 学习率调度器的 gamma 值。 + """ + super().__init__() + self.save_hyperparameters() + + # 输入投影层：将输入特征从 input_dim 映射到 d_model + self.input_projection = nn.Linear(input_dim, d_model) + + # Mamba2 模块 + self.mamba = Mamba2( + d_model=d_model, + d_state=d_state, + d_conv=d_conv, + expand=expand + ) + + # 池化层：对 Mamba2 的输出进行平均池化，得到固定长度的特征表示 + self.pool = nn.AdaptiveAvgPool1d(1) # 输出形状 [batch_size, d_model, 1] + + # 回归头：将池化后的特征映射到目标值 + self.fc1 = nn.Linear(d_model, d_model // 2) + self.dropout = nn.Dropout(dropout) + self.fc2 = nn.Linear(d_model // 2, 1) # 单输出 + + # 损失函数 + self.criterion = nn.MSELoss() + + # 评价指标 + self.train_mse = MeanMetric() + self.val_mse = MeanMetric() + self.test_mse = MeanMetric() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + 前向传播。 + + :param x: 输入张量，形状为 [batch_size, seq_len, input_dim] + :return: 预测张量，形状为 [batch_size] + """ + # 输入投影 + x = self.input_projection(x) # [batch_size, seq_len, d_model] + + # Mamba2 模块 + x = self.mamba(x) # [batch_size, seq_len, d_model] + + # 转置为 [batch_size, d_model, seq_len] 以适应 AdaptiveAvgPool1d + x = x.permute(0, 2, 1) # [batch_size, d_model, seq_len] + + # 池化 + x = self.pool(x) # [batch_size, d_model, 1] + x = x.squeeze(-1) # [batch_size, d_model] + + # 回归头 + x = F.relu(self.fc1(x)) # [batch_size, d_model // 2] + x = self.dropout(x) + x = self.fc2(x).squeeze(1) # [batch_size] + + return x + + def training_step(self, batch, batch_idx): + """ + 训练步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + :return: 损失值。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.train_mse(y_pred, y) + self.log("train/mse", self.train_mse, on_step=False, on_epoch=True, prog_bar=True) + + return loss + + def validation_step(self, batch, batch_idx): + """ + 验证步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.val_mse(y_pred, y) + self.log("val/mse", self.val_mse, on_step=False, on_epoch=True, prog_bar=True) + + def test_step(self, batch, batch_idx): + """ + 测试步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("test/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.test_mse(y_pred, y) + self.log("test/mse", self.test_mse, on_step=False, on_epoch=True, prog_bar=True) + + def configure_optimizers(self): + """ + 配置优化器和学习率调度器。 + + :return: 优化器和调度器配置。 + """ + optimizer = Adam(self.parameters(), lr=self.hparams.lr) + scheduler = StepLR( + optimizer, + step_size=self.hparams.scheduler_step_size, + gamma=self.hparams.scheduler_gamma, + ) + return [optimizer], [scheduler] diff --git a/src/models/PPGTransformerModule.py b/src/models/PPGTransformerModule.py index e69de29..07d86f1 100644 --- a/src/models/PPGTransformerModule.py +++ b/src/models/PPGTransformerModule.py @@ -0,0 +1,192 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from lightning import LightningModule +from torch.optim import Adam +from torch.optim.lr_scheduler import StepLR +from torchmetrics import MeanMetric + +class PositionalEncoding(nn.Module): + def __init__(self, d_model: int, max_len: int = 5000): + """ + 位置编码模块。 + + :param d_model: 输入特征的维度。 + :param max_len: 支持的最大序列长度。 + """ + super(PositionalEncoding, self).__init__() + + # 创建一个 [max_len, d_model] 的位置编码矩阵 + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # [max_len, 1] + div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # [d_model/2] + + pe[:, 0::2] = torch.sin(position * div_term) # 偶数维度 + pe[:, 1::2] = torch.cos(position * div_term) # 奇数维度 + + pe = pe.unsqueeze(0) # [1, max_len, d_model] + self.register_buffer('pe', pe) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + 为输入添加位置编码。 + + :param x: 输入张量，形状为 [batch_size, seq_len, d_model] + :return: 添加位置编码后的张量，形状相同 + """ + x = x + self.pe[:, :x.size(1), :] + return x + +class PPGTransformerModule(LightningModule): + def __init__( + self, + input_dim: int = 14, # 每个时间步的特征数 + d_model: int = 128, # Transformer模型的特征维度 + nhead: int = 8, # 多头注意力机制的头数 + num_encoder_layers: int = 4, # Transformer编码器层数 + dim_feedforward: int = 256, # 前馈网络的维度 + dropout: float = 0.1, # Dropout比例 + lr: float = 1e-3, # 学习率 + scheduler_step_size: int = 10, # 学习率调度器的步长 + scheduler_gamma: float = 0.1, # 学习率调度器的gamma值 + ): + """ + LightningModule 使用 Transformer 进行 PPG 数据预测。 + + :param input_dim: 每个时间步的输入特征数。 + :param d_model: Transformer模型的特征维度。 + :param nhead: 多头注意力机制的头数。 + :param num_encoder_layers: Transformer编码器层数。 + :param dim_feedforward: 前馈网络的维度。 + :param dropout: Dropout比例。 + :param lr: 学习率。 + :param scheduler_step_size: 学习率调度器的步长。 + :param scheduler_gamma: 学习率调度器的gamma值。 + """ + super().__init__() + self.save_hyperparameters() + + # 线性层将输入特征映射到d_model维度 + self.input_projection = nn.Linear(input_dim, d_model) + + # 位置编码 + self.positional_encoding = PositionalEncoding(d_model) + + # Transformer编码器 + encoder_layer = nn.TransformerEncoderLayer( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation='relu', + batch_first=True # 使用batch_first=True以匹配输入形状 [batch_size, seq_len, d_model] + ) + self.transformer_encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=num_encoder_layers + ) + + # 全连接层用于回归 + self.fc1 = nn.Linear(d_model, d_model // 2) + self.dropout = nn.Dropout(dropout) + self.fc2 = nn.Linear(d_model // 2, 1) # 单输出 + + # 损失函数 + self.criterion = nn.MSELoss() + + # 评价指标 + self.train_mse = MeanMetric() + self.val_mse = MeanMetric() + self.test_mse = MeanMetric() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + 前向传播。 + + :param x: 输入张量，形状为 [batch_size, seq_len, input_dim] + :return: 预测张量，形状为 [batch_size] + """ + # 线性映射 + x = self.input_projection(x) # [batch_size, seq_len, d_model] + + # 添加位置编码 + x = self.positional_encoding(x) # [batch_size, seq_len, d_model] + + # Transformer编码器 + x = self.transformer_encoder(x) # [batch_size, seq_len, d_model] + + # 池化操作：可以使用平均池化或取最后一个时间步 + # 这里使用平均池化 + x = x.mean(dim=1) # [batch_size, d_model] + + # 全连接层 + x = F.relu(self.fc1(x)) # [batch_size, d_model // 2] + x = self.dropout(x) + x = self.fc2(x).squeeze(1) # [batch_size] + + return x + + def training_step(self, batch, batch_idx): + """ + 训练步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + :return: 损失值。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.train_mse(y_pred, y) + self.log("train/mse", self.train_mse, on_step=False, on_epoch=True, prog_bar=True) + + return loss + + def validation_step(self, batch, batch_idx): + """ + 验证步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.val_mse(y_pred, y) + self.log("val/mse", self.val_mse, on_step=False, on_epoch=True, prog_bar=True) + + def test_step(self, batch, batch_idx): + """ + 测试步骤。 + + :param batch: 数据批次。 + :param batch_idx: 批次索引。 + """ + x, y = batch # x: [batch_size, window_size, 14], y: [batch_size] + y_pred = self.forward(x) # [batch_size] + loss = self.criterion(y_pred, y) + + # 记录损失和指标 + self.log("test/loss", loss, on_step=False, on_epoch=True, prog_bar=True) + self.test_mse(y_pred, y) + self.log("test/mse", self.test_mse, on_step=False, on_epoch=True, prog_bar=True) + + def configure_optimizers(self): + """ + 配置优化器和学习率调度器。 + + :return: 优化器和调度器配置。 + """ + optimizer = Adam(self.parameters(), lr=self.hparams.lr) + scheduler = StepLR( + optimizer, + step_size=self.hparams.scheduler_step_size, + gamma=self.hparams.scheduler_gamma, + ) + return [optimizer], [scheduler]