16 KiB
16 KiB
02 — 模型训练与 ONNX 导出详解
一、建模思路
1.1 为什么用 1D-CNN 而不是 LSTM？
电力波形是时序信号，直觉上 LSTM/GRU 更"语义匹配"。但部署到 NPU 时：
| 维度 | 1D-CNN | LSTM/BiLSTM | Transformer |
|---|---|---|---|
| NPU 原生支持 | ✅ 完美 | ⚠️ 部分 | ⚠️ 部分(RKNN2.0+) |
| INT8 量化精度损失 | <0.5% | 3~10% | 2~5% |
| 推理速度 (相同参数量) | 最快 | 慢 5~10× | 慢 3~8× |
| 感受野 | 可调(kernel_size/dilation) | 全序列 | 全序列 |
| 时域特征提取能力 | 局部→全局 | 顺序依赖 | 自注意力 |
结论：1D-CNN 通过堆叠层数 + 空洞卷积(dilation) 可以获得足够大的感受野，同时在 NPU 上有最优的硬件加速。电力故障的特征（过零点突变、幅值骤变、谐波畸变）本质上都是局部波形形态变化，CNN 完全能捕获。
1.2 建议的网络结构
输入: [batch, 1, 128] (1通道电流/电压, 128个采样点, 约2个周波@1.6kHz)
┌─────────────────────────────────────────┐
│ Conv1D(1→16, k=7, s=1, p=same) │ 浅层: 大卷积核捕获局部突变
│ BatchNorm1D + ReLU │
│ MaxPool1D(k=2) │ → [batch, 16, 64]
├─────────────────────────────────────────┤
│ Conv1D(16→32, k=5, s=1, p=same) │ 中层: 提取复合特征
│ BatchNorm1D + ReLU │
│ MaxPool1D(k=2) │ → [batch, 32, 32]
├─────────────────────────────────────────┤
│ Conv1D(32→64, k=3, s=1, p=same) │ 深层: 细粒度特征
│ BatchNorm1D + ReLU │
│ AdaptiveAvgPool1d(1) │ → [batch, 64, 1]
├─────────────────────────────────────────┤
│ FC(64→num_classes) │ 分类头
│ 输出: [batch, num_classes] │
└─────────────────────────────────────────┘
总参数量: 约 25K，模型大小 < 100KB (FP32), < 30KB (INT8)
1.3 模型设计原则（面向 NPU 部署）
- 尽量用 ReLU，不要用 LeakyReLU/PReLU（NPU 支持不完善）
- 避免动态 shape，输入尺寸固定（如
(1, 128)） - BatchNorm 不删，RKNN 转换时会自动融合 BN → Conv（减少推理量）
- 不用 Dropout（推理时本来就不生效，还占计算图节点）
- AdaptiveAvgPool1d 可以用，NPU 原生支持全局平均池化
二、数据准备
2.1 故障类型定义
电力系统常见的短路故障分类（以 5 分类为例）：
| 标签 | 故障类型 | 英文 | 波形特征 |
|---|---|---|---|
| 0 | 正常 | Normal | 标准正弦波 |
| 1 | A相单相接地 | A-G Fault | A相电压骤降～0, 零序电压升高 |
| 2 | BC相间短路 | BC Fault | B/C相电流激增 5~20×, 电压跌落 |
| 3 | AB两相接地 | AB-G Fault | 两相电压跌落 + 零序分量 |
| 4 | 三相短路 | ABC Fault | 三相电流对称激增, 电压严重跌落 |
如果需要更细的分类，可以扩展到 10 分类（加入高阻接地、电弧故障、铁磁谐振等）。
2.2 数据预处理流程
COMTRADE/CSV 录波文件
│
▼
┌──────────────────────────────────────┐
│ 1. 按通道提取 (A相电流 / A相电压 / 3I0) │
│ 每个通道独立训练一个模型 │
│ 或者多通道拼成 input_ch=3 │
├──────────────────────────────────────┤
│ 2. 滑动窗口切割 │
│ window = 128 点 │
│ stride = 64 点 (50% 重叠) │
│ 每个故障录波 → N 个样本 │
├──────────────────────────────────────┤
│ 3. 归一化 │
│ 方式A (推荐): Z-score │
│ x_norm = (x - μ) / σ │
│ μ,σ 从训练集统计, 写入配置文件 │
│ 方式B: Min-Max (仅用于特定场景) │
│ x_norm = (x - x_min)/(x_max - x_min)│
├──────────────────────────────────────┤
│ 4. 标签编码 │
│ one-hot [0,0,1,0,0] → class=2 │
├──────────────────────────────────────┤
│ 5. 数据集划分 (7:2:1) │
│ 训练集: 70% (实际训练) │
│ 验证集: 20% (调超参数) │
│ 测试集: 10% (仅最终评估) │
└──────────────────────────────────────┘
2.3 数据增强（解决不平衡问题）
电力故障数据天然不均衡——正常运行时间远大于故障时间。
# 对少数类做时间域增强
def augment_waveform(x):
# 1. 轻微时间偏移 (等效于改变故障起始角)
shift = np.random.randint(-8, 8)
x = np.roll(x, shift)
# 2. 幅度微扰 (±5%)
amp_noise = 1.0 + np.random.uniform(-0.05, 0.05)
x = x * amp_noise
# 3. 高斯噪声 (模拟传感器噪声)
noise = np.random.normal(0, 0.01, x.shape)
x = x + noise
return x
三、完整训练脚本
code/train_1dcnn.py：
"""
电力波形故障识别 — 1D-CNN 训练脚本
目标: 训练一个轻量级 1D-CNN，导出 ONNX，然后转 RKNN 部署到 RK3568 NPU
"""
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import os
# ============================================================
# 1. 模型定义
# ============================================================
class Fault1DCNN(nn.Module):
"""
轻量级 1D-CNN，面向 NPU 部署
参数量 ~25K，模型 FP32 大小 ~100KB
"""
def __init__(self, num_classes=5, input_channels=1, seq_len=128):
super().__init__()
# 第1卷积块: 大卷积核捕获局部突变 (过零点、尖峰)
self.block1 = nn.Sequential(
nn.Conv1d(input_channels, 16, kernel_size=7,
stride=1, padding='same', bias=False),
nn.BatchNorm1d(16),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=2), # 128 → 64
)
# 第2卷积块: 中等卷积核提取复合特征
self.block2 = nn.Sequential(
nn.Conv1d(16, 32, kernel_size=5,
stride=1, padding='same', bias=False),
nn.BatchNorm1d(32),
nn.ReLU(inplace=True),
nn.MaxPool1d(kernel_size=2), # 64 → 32
)
# 第3卷积块: 小卷积核细粒度特征
self.block3 = nn.Sequential(
nn.Conv1d(32, 64, kernel_size=3,
stride=1, padding='same', bias=False),
nn.BatchNorm1d(64),
nn.ReLU(inplace=True),
)
# 全局池化 → 固定长度特征向量
self.gap = nn.AdaptiveAvgPool1d(1) # 32 → 1
# 分类头
self.classifier = nn.Linear(64, num_classes)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv1d):
nn.init.kaiming_normal_(m.weight, mode='fan_out')
elif isinstance(m, nn.BatchNorm1d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
def forward(self, x):
# x: (batch, 1, 128)
x = self.block1(x) # (batch, 16, 64)
x = self.block2(x) # (batch, 32, 32)
x = self.block3(x) # (batch, 64, 32)
x = self.gap(x) # (batch, 64, 1)
x = x.view(x.size(0), -1) # (batch, 64)
x = self.classifier(x) # (batch, num_classes)
return x
# ============================================================
# 2. 数据集
# ============================================================
class WaveformDataset(Dataset):
"""
从 numpy 文件加载波形数据集
预期数据格式: data.npy (N, 128) labels.npy (N,)
或自定义 loader 从 COMTRADE 文件加载
"""
def __init__(self, data_path, label_path, augment=False):
self.data = np.load(data_path).astype(np.float32)
self.labels = np.load(label_path).astype(np.int64)
self.augment = augment
# 确保是 (N, 1, 128) 格式
if self.data.ndim == 2:
self.data = self.data[:, np.newaxis, :]
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
x = self.data[idx]
y = self.labels[idx]
if self.augment:
x = self._augment(x)
return torch.from_numpy(x), torch.tensor(y)
def _augment(self, x):
"""时间域数据增强 (对 numpy array 操作)"""
x = x.copy()
# 随机时移 ±8 点
shift = np.random.randint(-8, 9)
if shift != 0:
x = np.roll(x, shift, axis=-1)
# 幅度微扰 ±3%
x *= (1.0 + np.random.uniform(-0.03, 0.03))
# 加性噪声
x += np.random.normal(0, 1e-3, x.shape)
return x
# ============================================================
# 3. 训练
# ============================================================
def train_one_epoch(model, loader, criterion, optimizer, device):
model.train()
total_loss, correct, total = 0.0, 0, 0
for x, y in loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
total_loss += loss.item() * x.size(0)
correct += (logits.argmax(1) == y).sum().item()
total += x.size(0)
return total_loss / total, correct / total
def validate(model, loader, criterion, device):
model.eval()
total_loss, correct, total = 0.0, 0, 0
all_preds, all_labels = [], []
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
logits = model(x)
loss = criterion(logits, y)
total_loss += loss.item() * x.size(0)
preds = logits.argmax(1)
correct += (preds == y).sum().item()
total += x.size(0)
all_preds.extend(preds.cpu().numpy())
all_labels.extend(y.cpu().numpy())
return (total_loss / total, correct / total,
all_preds, all_labels)
def main():
# ---- 配置 ----
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 64
NUM_EPOCHS = 100
LR = 0.001
NUM_CLASSES = 5
SEQ_LEN = 128
INPUT_CH = 1
# ---- 数据加载 ----
train_ds = WaveformDataset('data/train_data.npy',
'data/train_labels.npy', augment=True)
val_ds = WaveformDataset('data/val_data.npy',
'data/val_labels.npy', augment=False)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE,
shuffle=True, num_workers=2)
val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE,
shuffle=False, num_workers=2)
print(f"训练集: {len(train_ds)}, 验证集: {len(val_ds)}")
# ---- 模型 ----
model = Fault1DCNN(num_classes=NUM_CLASSES,
input_channels=INPUT_CH, seq_len=SEQ_LEN)
model = model.to(DEVICE)
print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")
# ---- 优化器 & 损失函数 ----
optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='max', factor=0.5, patience=10, verbose=True)
criterion = nn.CrossEntropyLoss()
# ---- 训练循环 ----
best_acc = 0.0
for epoch in range(1, NUM_EPOCHS + 1):
train_loss, train_acc = train_one_epoch(
model, train_loader, criterion, optimizer, DEVICE)
val_loss, val_acc, preds, labels = validate(
model, val_loader, criterion, DEVICE)
scheduler.step(val_acc)
if epoch % 10 == 0 or epoch == 1:
print(f"Epoch {epoch:3d} | "
f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2%} | "
f"Val Loss: {val_loss:.4f} Acc: {val_acc:.2%}")
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'checkpoints/best_model.pth')
print(f"\n最佳验证准确率: {best_acc:.2%}")
# ---- 最终评估 ----
print("\n" + "="*50)
print("最终验证集分类报告:")
print(classification_report(labels, preds,
target_names=['Normal','A-G','BC','AB-G','ABC']))
print("="*50)
# ---- 导出 ONNX ----
export_onnx(model, 'checkpoints/best_model.pth', SEQ_LEN, INPUT_CH, DEVICE)
# ============================================================
# 4. ONNX 导出 (关键步骤，影响后续 RKNN 转换成功率)
# ============================================================
def export_onnx(model, checkpoint_path, seq_len, input_ch, device):
"""导出 ONNX，注意算子兼容性"""
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
model.eval()
dummy_input = torch.randn(1, input_ch, seq_len).to(device)
output_path = 'checkpoints/fault_1dcnn.onnx'
torch.onnx.export(
model,
dummy_input,
output_path,
input_names=['waveform'],
output_names=['logits'],
dynamic_axes=None, # ⚠️ 不要导出动态 shape!
opset_version=12, # RKNN 推荐 opset 11~12
do_constant_folding=True, # 常量折叠, 减少计算图节点
export_params=True,
verbose=False,
)
print(f"\nONNX 模型已导出: {output_path}")
# 验证 ONNX
import onnx
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX 模型验证通过 ✓")
# 打印计算图算子统计（确认没有不支持的算子）
op_types = [n.op_type for n in onnx_model.graph.node]
from collections import Counter
op_count = Counter(op_types)
print(f"算子统计: {dict(op_count)}")
# 预期: {'Conv':3, 'Relu':3, 'MaxPool':2, 'GlobalAveragePool':1,
# 'Gemm':1, 'BatchNorm':3} (BN 后续会被融合)
if __name__ == '__main__':
os.makedirs('checkpoints', exist_ok=True)
main()
3.2 训练建议
# 安装依赖
pip install torch numpy scikit-learn onnx onnxruntime
# 准备数据
# data/train_data.npy (N_train, 128) float32
# data/train_labels.npy (N_train,) int64
# data/val_data.npy (N_val, 128)
# data/val_labels.npy (N_val,)
# 训练
python train_1dcnn.py
# 验证 ONNX 推理结果与 PyTorch 一致
python -c "
import onnxruntime as ort
import numpy as np
import torch
# PyTorch 推理
model = ... # 加载训练好的模型
dummy = torch.randn(1,1,128)
pt_out = model(dummy).detach().numpy()
# ONNX 推理
sess = ort.InferenceSession('checkpoints/fault_1dcnn.onnx')
onnx_out = sess.run(None, {'waveform': dummy.numpy()})[0]
# 比较
print('最大误差:', np.max(np.abs(pt_out - onnx_out)))
# 期望: < 1e-5
"
四、训练要点小结
| 要点 | 做法 | 原因 |
|---|---|---|
| 网络选型 | 1D-CNN, 不要 LSTM/Transformer | NPU 的 CNN 算子最优 |
| 激活函数 | 只用 ReLU | NPU 原生支持，量化稳定 |
| 卷积核设计 | 浅层大核(7) → 深层小核(3) | 先捕获突变，再提取细节 |
| 输入固定 | shape=(1,1,128)，不导出动态轴 | NPU 不支持动态 shape |
| ONNX opset | 11 或 12 | 太高可能有不兼容算子 |
| 检查 ONNX 算子 | export 后打印 op_type 列表 | 提前发现不支持算子 |
| 训练轮次 | 100~200 epoch, ReduceLROnPlateau | 小模型收敛快，不需太多轮 |