318 lines
10 KiB
Python
318 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
电力波形故障识别 — 1D-CNN 训练脚本 (完整版)
|
||
|
||
使用方法:
|
||
python train_1dcnn.py --data_dir ./data --epochs 100
|
||
|
||
数据准备:
|
||
data/
|
||
train_data.npy (N_train, 128) float32 训练样本
|
||
train_labels.npy (N_train,) int64 训练标签
|
||
val_data.npy (N_val, 128)
|
||
val_labels.npy (N_val,)
|
||
test_data.npy (N_test, 128)
|
||
test_labels.npy (N_test,)
|
||
|
||
输出:
|
||
checkpoints/best_model.pth 最佳权重
|
||
checkpoints/fault_1dcnn.onnx ONNX 模型 (用于 RKNN 转换)
|
||
"""
|
||
|
||
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
|
||
import os
|
||
import argparse
|
||
from collections import Counter
|
||
|
||
|
||
# ============================================================
|
||
# 故障类型名称
|
||
# ============================================================
|
||
FAULT_NAMES = [
|
||
'Normal', # 0 正常
|
||
'A-G', # 1 A相单相接地
|
||
'BC', # 2 BC相间短路
|
||
'AB-G', # 3 AB两相接地
|
||
'ABC', # 4 三相短路
|
||
]
|
||
|
||
NUM_CLASSES = len(FAULT_NAMES)
|
||
SEQ_LEN = 128
|
||
INPUT_CH = 1
|
||
|
||
|
||
# ============================================================
|
||
# 模型定义
|
||
# ============================================================
|
||
class Fault1DCNN(nn.Module):
|
||
"""
|
||
轻量级 1D-CNN，面向 RK3568 NPU 部署
|
||
参数量 ~25K，模型 FP32 大小 ~100KB
|
||
"""
|
||
|
||
def __init__(self, num_classes=NUM_CLASSES,
|
||
input_channels=INPUT_CH, seq_len=SEQ_LEN):
|
||
super().__init__()
|
||
_ = seq_len # unused, kept for API compatibility
|
||
|
||
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
|
||
)
|
||
|
||
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
|
||
)
|
||
|
||
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 = self.block1(x) # (B, 16, 64)
|
||
x = self.block2(x) # (B, 32, 32)
|
||
x = self.block3(x) # (B, 64, 32)
|
||
x = self.gap(x) # (B, 64, 1)
|
||
x = x.view(x.size(0), -1)
|
||
x = self.classifier(x) # (B, num_classes)
|
||
return x
|
||
|
||
|
||
# ============================================================
|
||
# 数据集 (含数据增强)
|
||
# ============================================================
|
||
class WaveformDataset(Dataset):
|
||
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
|
||
|
||
if self.data.ndim == 2:
|
||
self.data = self.data[:, np.newaxis, :]
|
||
|
||
# 统计归一化参数
|
||
self.mean = self.data.mean()
|
||
self.std = self.data.std() + 1e-8
|
||
|
||
def __len__(self):
|
||
return len(self.data)
|
||
|
||
def __getitem__(self, idx):
|
||
x = self.data[idx].copy()
|
||
y = self.labels[idx]
|
||
|
||
# Z-score 归一化
|
||
x = (x - self.mean) / self.std
|
||
|
||
if self.augment:
|
||
x = self._augment(x)
|
||
|
||
return torch.from_numpy(x), torch.tensor(y)
|
||
|
||
@staticmethod
|
||
def _augment(x):
|
||
"""时间域数据增强 (numpy, shape=(1, L))"""
|
||
shift = np.random.randint(-8, 9)
|
||
if shift != 0:
|
||
x = np.roll(x, shift, axis=-1)
|
||
x *= (1.0 + np.random.uniform(-0.03, 0.03))
|
||
x += np.random.normal(0, 1e-3, x.shape).astype(np.float32)
|
||
return x
|
||
|
||
|
||
# ============================================================
|
||
# 训练 / 验证
|
||
# ============================================================
|
||
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()
|
||
loss = criterion(model(x), y)
|
||
loss.backward()
|
||
optimizer.step()
|
||
total_loss += loss.item() * x.size(0)
|
||
correct += (model(x).argmax(1) == y).sum().item() # re-run for stats
|
||
total += x.size(0)
|
||
|
||
# Recalculate accuracy properly
|
||
model.eval()
|
||
correct, total = 0, 0
|
||
with torch.no_grad():
|
||
for x, y in loader:
|
||
x, y = x.to(device), y.to(device)
|
||
logits = model(x)
|
||
correct += (logits.argmax(1) == y).sum().item()
|
||
total += x.size(0)
|
||
|
||
return total_loss / total, correct / total
|
||
|
||
|
||
@torch.no_grad()
|
||
def validate(model, loader, criterion, device):
|
||
model.eval()
|
||
total_loss, correct, total = 0.0, 0, 0
|
||
all_preds, all_labels = [], []
|
||
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
|
||
|
||
|
||
# ============================================================
|
||
# ONNX 导出
|
||
# ============================================================
|
||
def export_onnx(model, checkpoint_path, output_path, device):
|
||
state = torch.load(checkpoint_path, map_location=device)
|
||
model.load_state_dict(state)
|
||
model.eval()
|
||
|
||
dummy = torch.randn(1, INPUT_CH, SEQ_LEN).to(device)
|
||
|
||
torch.onnx.export(
|
||
model, dummy, output_path,
|
||
input_names=['waveform'],
|
||
output_names=['logits'],
|
||
dynamic_axes=None,
|
||
opset_version=12,
|
||
do_constant_folding=True,
|
||
export_params=True,
|
||
verbose=False,
|
||
)
|
||
|
||
# 验证 ONNX
|
||
try:
|
||
import onnx
|
||
onnx_model = onnx.load(output_path)
|
||
onnx.checker.check_model(onnx_model)
|
||
op_types = [n.op_type for n in onnx_model.graph.node]
|
||
print(f"\nONNX 算子统计: {dict(Counter(op_types))}")
|
||
print("ONNX 模型验证通过 ✓")
|
||
except ImportError:
|
||
print("WARNING: 未安装 onnx 库，跳过 ONNX 验证")
|
||
|
||
print(f"ONNX 模型已导出: {output_path}")
|
||
|
||
|
||
# ============================================================
|
||
# Main
|
||
# ============================================================
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='电力波形故障识别 1D-CNN')
|
||
parser.add_argument('--data_dir', default='data', help='数据目录')
|
||
parser.add_argument('--batch_size', type=int, default=64)
|
||
parser.add_argument('--epochs', type=int, default=100)
|
||
parser.add_argument('--lr', type=float, default=0.001)
|
||
parser.add_argument('--device', default='auto')
|
||
args = parser.parse_args()
|
||
|
||
# 设备
|
||
if args.device == 'auto':
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
else:
|
||
device = torch.device(args.device)
|
||
print(f"设备: {device}")
|
||
|
||
# 数据
|
||
train_ds = WaveformDataset(
|
||
f'{args.data_dir}/train_data.npy',
|
||
f'{args.data_dir}/train_labels.npy', augment=True)
|
||
val_ds = WaveformDataset(
|
||
f'{args.data_dir}/val_data.npy',
|
||
f'{args.data_dir}/val_labels.npy', augment=False)
|
||
|
||
train_loader = DataLoader(train_ds, batch_size=args.batch_size,
|
||
shuffle=True, num_workers=2)
|
||
val_loader = DataLoader(val_ds, batch_size=args.batch_size,
|
||
shuffle=False, num_workers=2)
|
||
|
||
print(f"训练集: {len(train_ds):,} 验证集: {len(val_ds):,}")
|
||
print(f"归一化: mean={train_ds.mean:.3f}, std={train_ds.std:.3f}")
|
||
|
||
# 数据不平衡检测
|
||
label_counts = Counter(train_ds.labels)
|
||
print(f"训练集类别分布: {dict(label_counts)}")
|
||
|
||
# 模型
|
||
model = Fault1DCNN().to(device)
|
||
n_params = sum(p.numel() for p in model.parameters())
|
||
print(f"模型参数量: {n_params:,} (~{n_params*4/1024:.0f} KB FP32)")
|
||
|
||
# 训练
|
||
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)
|
||
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
|
||
optimizer, mode='max', factor=0.5, patience=10, verbose=True)
|
||
criterion = nn.CrossEntropyLoss()
|
||
|
||
os.makedirs('checkpoints', exist_ok=True)
|
||
best_acc = 0.0
|
||
|
||
for epoch in range(1, args.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" + "=" * 55)
|
||
print(classification_report(labels, preds, target_names=FAULT_NAMES))
|
||
print("=" * 55)
|
||
|
||
# 导出 ONNX
|
||
export_onnx(model, 'checkpoints/best_model.pth',
|
||
'checkpoints/fault_1dcnn.onnx', device)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|