手把手复现SAM-BiGRU锂电池寿命预测模型（附PyTorch 2.1+代码与CALCE数据集）

手把手复现SAM-BiGRU锂电池寿命预测模型附PyTorch 2.1代码与CALCE数据集锂电池作为新能源领域的核心储能元件其剩余使用寿命RUL预测一直是工业界和学术界关注的焦点。今天我们将从零开始完整复现一篇基于自注意力机制SAM和双向门控循环单元BiGRU的锂电池寿命预测模型。不同于简单的代码搬运本文将深入每个技术细节包括数据预处理的坑点、模型架构的PyTorch实现技巧以及如何调参才能达到论文中的高精度预测效果。1. 环境准备与数据加载1.1 PyTorch环境配置推荐使用conda创建虚拟环境避免包冲突conda create -n sam_bigru python3.9 conda activate sam_bigru pip install torch2.1.0 torchvision torchaudio pip install pandas numpy matplotlib scikit-learn注意PyTorch 2.1版本对混合精度训练有更好的支持这对BiGRU这类序列模型尤为重要1.2 CALCE数据集解析CALCE数据集包含多个电池的充放电循环数据关键字段包括字段名说明单位典型值Cycle循环次数-1~500Voltage端电压V2.7~4.2Current充放电电流A±1.1Capacity当前容量Ah1.1~0.7数据加载的核心代码import pandas as pd def load_calce_data(battery_id): df pd.read_csv(fCALCE_CS2_{battery_id}.csv) # 关键预处理计算每个循环的容量衰减率 df[Capacity_Degradation] 1 - df[Capacity] / df[Capacity].max() return df[[Cycle, Capacity_Degradation]].values2. 数据预处理关键技术2.1 滑动窗口异常值处理锂电池数据常因测量误差出现异常点采用动态阈值检测def sliding_window_outlier_detection(data, window_size10, threshold3): cleaned_data data.copy() for i in range(len(data)): window_start max(0, i - window_size) window data[window_start:i] mean, std window.mean(), window.std() if abs(data[i] - mean) threshold * std: # 使用前后窗口的中位数插值 cleaned_data[i] np.median(window) return cleaned_data2.2 序列标准化与分割采用电池专属的归一化方式避免不同电池间的量纲影响from sklearn.preprocessing import MinMaxScaler def prepare_sequences(data, seq_length30): scaler MinMaxScaler(feature_range(0, 1)) scaled_data scaler.fit_transform(data.reshape(-1, 1)) X, y [], [] for i in range(len(scaled_data) - seq_length): X.append(scaled_data[i:iseq_length]) y.append(scaled_data[iseq_length]) return np.array(X), np.array(y), scaler3. SAM-BiGRU模型架构实现3.1 自注意力模块代码详解import torch.nn as nn class SelfAttention(nn.Module): def __init__(self, embed_size): super(SelfAttention, self).__init__() self.query nn.Linear(embed_size, embed_size) self.key nn.Linear(embed_size, embed_size) self.value nn.Linear(embed_size, embed_size) self.softmax nn.Softmax(dim-1) def forward(self, x): # x shape: (batch, seq_len, embed_size) Q self.query(x) K self.key(x) V self.value(x) attention_scores torch.matmul(Q, K.transpose(1, 2)) / (x.size(-1) ** 0.5) attention self.softmax(attention_scores) out torch.matmul(attention, V) return out3.2 BiGRU与模型整合class SAMBiGRU(nn.Module): def __init__(self, input_size1, hidden_size64, output_size1): super(SAMBiGRU, self).__init__() self.attention SelfAttention(input_size) self.bigru nn.GRU(input_size, hidden_size, bidirectionalTrue, batch_firstTrue) self.fc nn.Linear(hidden_size * 2, output_size) def forward(self, x): att_out self.attention(x) gru_out, _ self.bigru(att_out) # 取最后一个时间步的输出 out self.fc(gru_out[:, -1, :]) return out4. 模型训练与调优实战4.1 训练参数配置from torch.optim import Adam from torch.utils.data import DataLoader, TensorDataset # 超参数设置 params { seq_length: 30, # 输入序列长度 batch_size: 32, # 批大小 hidden_size: 64, # BiGRU隐藏层维度 learning_rate: 0.001, epochs: 200, patience: 15 # 早停等待轮次 } # 初始化模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model SAMBiGRU(hidden_sizeparams[hidden_size]).to(device) optimizer Adam(model.parameters(), lrparams[learning_rate]) criterion nn.MSELoss()4.2 训练循环关键代码def train_model(model, train_loader, val_loader, params): best_loss float(inf) patience_counter 0 for epoch in range(params[epochs]): model.train() train_loss 0 for X_batch, y_batch in train_loader: optimizer.zero_grad() outputs model(X_batch) loss criterion(outputs, y_batch) loss.backward() optimizer.step() train_loss loss.item() # 验证阶段 val_loss evaluate(model, val_loader, criterion) # 早停机制 if val_loss best_loss: best_loss val_loss patience_counter 0 torch.save(model.state_dict(), best_model.pth) else: patience_counter 1 if patience_counter params[patience]: print(fEarly stopping at epoch {epoch}) break4.3 提升预测精度的三个技巧学习率预热前5个epoch使用线性增长的学习率def adjust_learning_rate(optimizer, epoch, warmup5, init_lr1e-5, max_lr1e-3): if epoch warmup: lr init_lr (max_lr - init_lr) * epoch / warmup for param_group in optimizer.param_groups: param_group[lr] lr混合精度训练减少显存占用加快训练速度from torch.cuda.amp import GradScaler, autocast scaler GradScaler() with autocast(): outputs model(X_batch) loss criterion(outputs, y_batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()多尺度特征提取在BiGRU前增加不同尺度的CNN层class MultiScaleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv1d(1, 16, kernel_size3, padding1) self.conv2 nn.Conv1d(1, 16, kernel_size5, padding2) def forward(self, x): x1 F.relu(self.conv1(x)) x2 F.relu(self.conv2(x)) return torch.cat([x1, x2], dim1)5. 结果评估与工业部署建议5.1 评估指标实现from sklearn.metrics import mean_absolute_error, mean_squared_error def evaluate_rul(y_true, y_pred): mae mean_absolute_error(y_true, y_pred) rmse np.sqrt(mean_squared_error(y_true, y_pred)) error_percent np.mean(np.abs(y_true - y_pred) / y_true) * 100 print(fMAE: {mae:.4f}, RMSE: {rmse:.4f}, Error%: {error_percent:.2f}%) return {MAE: mae, RMSE: rmse, Error%: error_percent}5.2 实际部署优化策略模型轻量化使用知识蒸馏训练小型BiGRUdef distillation_loss(student_output, teacher_output, temp2.0): return F.kl_div( F.log_softmax(student_output/temp, dim1), F.softmax(teacher_output/temp, dim1), reductionbatchmean) * (temp**2)边缘计算适配转换为ONNX格式torch.onnx.export(model, dummy_input, sam_bigru.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})持续学习机制实现模型在线更新class OnlineLearner: def __init__(self, model, buffer_size100): self.model model self.buffer deque(maxlenbuffer_size) def update(self, new_data): self.buffer.append(new_data) if len(self.buffer) self.buffer.maxlen: self._fine_tune() def _fine_tune(self): # 小批量微调逻辑 pass在完成模型训练后我发现一个实用技巧将滑动窗口大小设置为电池典型失效周期的1/3左右CALCE数据集中约30-40个循环时模型对早期衰退信号的捕捉最为敏感。另外使用指数衰减的学习率调度gamma0.95比阶跃式衰减能获得更稳定的收敛效果。