红外火情时序预判 CNN-LSTM 模型
基于 CNN-LSTM 的阴燃蔓延时序预判实现从 “事后识别火情” 升级为 “提前预判高温扩张风险”。整套流程分为视频时序样本构建、CNN-LSTM 模型训练、ONNX 模型导出、嵌入式部署四大环节。数据来源红外摄像头录制1m 30s 分段 AVI 火情视频样本构造滑动窗口截取连续 8 帧红外图像为输入序列提取未来帧高温特征作为回归标签划分安全 / 阴燃扩张 / 明火三类风险标签标签提取对每组序列提取未来帧的高温像素面积、图像平均灰度、高温区域中心坐标作为回归标签同时标注 0 无扩张风险、1 阴燃持续扩张、2 即将出现明火三分类标签数据划分训练集 80%、验证集 20%统一归一化至 [0,1]适配模型输入。数据集规格共 8 个.npz时序文件数据集硬盘占用 12GB训练硬件笔记本 CPU内存有限无法一次性加载全部数据。初始代码一次性读取全部 npz 文件加载数组时 numpy 直接申请 4GB 以上内存程序初始化阶段直接崩溃MemoryError: Unable to allocate 4.17 GiB for an array with shape (559153152,) and data type float64摒弃全局一次性加载改为单文件分批训练不再扫描全部文件构建全局索引循环逐个读取 npz训练完成立刻关闭文件释放内存同一时间仅占用单个文件内存import os import torch import torch.nn as nn import torch.optim as optim import numpy as np from torch.utils.data import Dataset, DataLoader device torch.device(cuda if torch.cuda.is_available() else cpu) SEQ_LENGTH 8 H, W 192, 256 WEIGHT_SAVE temp_weight.pth class FireSingleFileDataset(Dataset): def __init__(self, file_path): print(fLoad single file: {os.path.basename(file_path)}) self.data np.load(file_path, mmap_moder) self.seq_arr self.data[seq] self.label_arr self.data[label] self.total len(self.seq_arr) print(fSamples in this file: {self.total}) def __len__(self): return self.total def __getitem__(self, idx): seq self.seq_arr[idx] label self.label_arr[idx] seq torch.from_numpy(seq).permute(0, 3, 1, 2).float() label torch.from_numpy(label).float() return seq, label def close(self): self.data.close() class ConvBlock(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.block nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.ReLU(), nn.MaxPool2d(2) ) def forward(self, x): return self.block(x) class CNNLSTM(nn.Module): def __init__(self): super().__init__() self.cnn nn.Sequential( ConvBlock(3, 16), ConvBlock(16, 32), ConvBlock(32, 64) ) self.lstm nn.LSTM(input_size64*24*32, hidden_size128, batch_firstTrue) self.reg_head nn.Linear(128, 4) self.cls_head nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 3)) def forward(self, x): b, seq, c, h, w x.shape feat_list [] for t in range(seq): feat self.cnn(x[:, t]) feat_list.append(feat.flatten(1)) seq_feat torch.stack(feat_list, dim1) lstm_out, _ self.lstm(seq_feat) last_feat lstm_out[:, -1, :] reg_out self.reg_head(last_feat) cls_out self.cls_head(last_feat) return reg_out, cls_out def train_one_file(model, file_path, epoch_num1): dataset FireSingleFileDataset(file_path) loader DataLoader(dataset, batch_size1, shuffleTrue, num_workers0) opt optim.Adam(model.parameters(), lr1e-4) loss_mse nn.MSELoss() loss_ce nn.CrossEntropyLoss() print(f Start train file {os.path.basename(file_path)} ) for e in range(epoch_num): total_loss 0.0 for seq, label in loader: seq, label seq.to(device), label.to(device) pred_reg, pred_cls model(seq) loss1 loss_mse(pred_reg, label) cls_label torch.clamp(torch.round(label[:, 0]).long(), 0, 2) loss2 loss_ce(pred_cls, cls_label) loss loss1 loss2 opt.zero_grad() loss.backward() opt.step() total_loss loss.item() avg_loss total_loss / len(loader) print(fEpoch {e1}/{epoch_num}, Loss: {avg_loss:.4f}) dataset.close() torch.save(model.state_dict(), WEIGHT_SAVE) print(fTemporary weight saved to {WEIGHT_SAVE}) if __name__ __main__: print( Start Initializing CNN-LSTM Model ) model CNNLSTM().to(device) print(fModel loaded on device: {device}) dataset_path rD:\图像处理\venv\train_seq all_files [os.path.join(dataset_path, f) for f in os.listdir(dataset_path) if f.endswith(.npz)] print(fTotal npz files: {len(all_files)}) if os.path.exists(WEIGHT_SAVE): print(fLoad previous weight from {WEIGHT_SAVE}) model.load_state_dict(torch.load(WEIGHT_SAVE, map_locationdevice)) for f in all_files: train_one_file(model, f, epoch_num1) print(Release memory, prepare next file\n) print( All files training complete, export ONNX Model ) model.eval() dummy torch.randn(1, SEQ_LENGTH, 3, H, W).to(device) save_name fire_predict.onnx full_save_path os.path.abspath(save_name) torch.onnx.export(model, dummy, full_save_path, opset_version12) print(fONNX model saved successfully! Path: {full_save_path}) print(Upload fire_predict.onnx to RK3568 /home via MobaXterm)每次仅打开单个 npz训练结束调用close()释放内存杜绝多文件同时占用内存每个文件训练完成保存临时权重下次运行可加载权重接续训练无需从头开始CNN 提取单帧红外热区空间特征LSTM 学习时序扩张规律同时输出热区回归预测值 火情风险三分类结果。