2. 边缘设备上的轴承故障分类

任务说明

在边缘设备（如STM32、ESP32或树莓派Pico）上部署一个轻量级神经网络，对轴承振动信号进行二分类：判断当前轴承状态为“正常”或“故障”。本案例聚焦于模型训练、量化与C代码导出，最终生成可直接移植到边缘设备的整数权重数组与推理函数。

实现方案

使用Python + NumPy从零构建一个两层全连接网络（输入层→隐藏层→输出层），不依赖TensorFlow/PyTorch等重型框架。输入为轴承振动信号的频域特征（取16个频段能量），输出为两类概率。训练完成后，将权重与偏置量化为int8整数，并生成标准C头文件，包含推理所需的全部常量与函数。

代码实现

import numpy as np

# ========== 1. 生成模拟轴承数据（正常/故障各200组） ==========
np.random.seed(42)
n_samples = 400
n_features = 16

# 正常轴承：低频段能量高，高频段能量低
normal = np.random.rand(200, n_features)
normal[:, :8] *= 0.8 + 0.2 * np.random.rand(200, 8)
normal[:, 8:] *= 0.1 + 0.1 * np.random.rand(200, 8)

# 故障轴承（如内圈裂纹）：高频段出现明显能量峰
fault = np.random.rand(200, n_features)
fault[:, :8] *= 0.2 + 0.1 * np.random.rand(200, 8)
fault[:, 8:] *= 0.7 + 0.3 * np.random.rand(200, 8)

X = np.vstack([normal, fault])
y = np.array([0]*200 + [1]*200).reshape(-1, 1)

# 归一化到[0,1]
X = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0) + 1e-8)

# 划分训练/测试集
idx = np.random.permutation(n_samples)
train_idx, test_idx = idx[:320], idx[320:]
X_train, X_test = X[train_idx], X[test_idx]
y_train, y_test = y[train_idx], y[test_idx]

# ========== 2. 两层全连接网络定义 ==========
input_dim = 16
hidden_dim = 8
output_dim = 1
lr = 0.5

# He初始化
W1 = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0/input_dim)
b1 = np.zeros((1, hidden_dim))
W2 = np.random.randn(hidden_dim, output_dim) * np.sqrt(2.0/hidden_dim)
b2 = np.zeros((1, output_dim))

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -30, 30)))

# ========== 3. 训练（100轮梯度下降） ==========
for epoch in range(100):
    # 前向传播
    z1 = X_train @ W1 + b1
    a1 = sigmoid(z1)
    z2 = a1 @ W2 + b2
    a2 = sigmoid(z2)
    
    # 损失（二分类交叉熵）
    loss = -np.mean(y_train * np.log(a2 + 1e-8) + (1-y_train) * np.log(1-a2 + 1e-8))
    
    # 反向传播
    dz2 = a2 - y_train
    dW2 = a1.T @ dz2 / len(X_train)
    db2 = np.sum(dz2, axis=0, keepdims=True) / len(X_train)
    
    da1 = dz2 @ W2.T
    dz1 = da1 * a1 * (1 - a1)
    dW1 = X_train.T @ dz1 / len(X_train)
    db1 = np.sum(dz1, axis=0, keepdims=True) / len(X_train)
    
    # 参数更新
    W1 -= lr * dW1
    b1 -= lr * db1
    W2 -= lr * dW2
    b2 -= lr * db2

# ========== 4. 测试精度 ==========
z1 = X_test @ W1 + b1
a1 = sigmoid(z1)
z2 = a1 @ W2 + b2
a2 = sigmoid(z2)
pred = (a2 > 0.5).astype(int)
acc = np.mean(pred == y_test)
print(f"测试集准确率: {acc:.3f}")

# ========== 5. 量化为int8并导出C头文件 ==========
def quantize(w, scale):
    q = np.round(w / scale).astype(np.int8)
    q = np.clip(q, -128, 127)
    return q

# 确定缩放因子（取最大值/127）
s1 = np.max(np.abs(W1)) / 127.0
s2 = np.max(np.abs(W2)) / 127.0
sb1 = np.max(np.abs(b1)) / 127.0
sb2 = np.max(np.abs(b2)) / 127.0

W1_q = quantize(W1, s1)
W2_q = quantize(W2, s2)
b1_q = quantize(b1, sb1)
b2_q = quantize(b2, sb2)

# 生成C头文件
with open("bearing_model.h", "w") as f:
    f.write("#ifndef BEARING_MODEL_H\n#define BEARING_MODEL_H\n\n")
    f.write("#include <stdint.h>\n\n")
    f.write(f"#define INPUT_DIM {input_dim}\n")
    f.write(f"#define HIDDEN_DIM {hidden_dim}\n")
    f.write(f"#define OUTPUT_DIM {output_dim}\n\n")
    
    f.write(f"// 缩放因子（float）\n")
    f.write(f"const float W1_SCALE = {s1:.8f}f;\n")
    f.write(f"const float W2_SCALE = {s2:.8f}f;\n")
    f.write(f"const float B1_SCALE = {sb1:.8f}f;\n")
    f.write(f"const float B2_SCALE = {sb2:.8f}f;\n\n")
    
    f.write("// 量化权重\n")
    f.write("const int8_t W1[INPUT_DIM][HIDDEN_DIM] = {\n")
    for i in range(input_dim):
        row = ", ".join(str(int(v)) for v in W1_q[i])
        f.write(f"    {{{row}}}{',' if i < input_dim-1 else ''}\n")
    f.write("};\n\n")
    
    f.write("const int8_t W2[HIDDEN_DIM][OUTPUT_DIM] = {\n")
    for i in range(hidden_dim):
        row = ", ".join(str(int(v)) for v in W2_q[i])
        f.write(f"    {{{row}}}{',' if i < hidden_dim-1 else ''}\n")
    f.write("};\n\n")
    
    f.write("const int8_t B1[HIDDEN_DIM] = {")
    f.write(", ".join(str(int(v)) for v in b1_q.flatten()))
    f.write("};\n\n")
    
    f.write("const int8_t B2[OUTPUT_DIM] = {")
    f.write(", ".join(str(int(v)) for v in b2_q.flatten()))
    f.write("};\n\n")
    
    f.write("// 推理函数：输入16个频段能量（已归一化到[0,1]），输出故障概率\n")
    f.write("static inline float bearing_inference(const float* input) {\n")
    f.write("    float hidden[HIDDEN_DIM];\n")
    f.write("    float out;\n")
    f.write("    \n")
    f.write("    // 隐藏层\n")
    f.write("    for (int j = 0; j < HIDDEN_DIM; j++) {\n")
    f.write("        float sum = (float)B1[j] * B1_SCALE;\n")
    f.write("        for (int i = 0; i < INPUT_DIM; i++) {\n")
    f.write("            sum += input[i] * ((float)W1[i][j] * W1_SCALE);\n")
    f.write("        }\n")
    f.write("        hidden[j] = 1.0f / (1.0f + expf(-sum));\n")
    f.write("    }\n")
    f.write("    \n")
    f.write("    // 输出层\n")
    f.write("    float sum = (float)B2[0] * B2_SCALE;\n")
    f.write("    for (int j = 0; j < HIDDEN_DIM; j++) {\n")
    f.write("        sum += hidden[j] * ((float)W2[j][0] * W2_SCALE);\n")
    f.write("    }\n")
    f.write("    out = 1.0f / (1.0f + expf(-sum));\n")
    f.write("    return out;  // 返回故障概率，>0.5判定为故障\n")
    f.write("}\n\n")
    f.write("#endif // BEARING_MODEL_H\n")

print("已生成 bearing_model.h")

落地与验证

运行环境：Python 3.8+，仅需NumPy库（pip install numpy）。

执行步骤：

1. 将上述代码保存为 train_bearing.py 并运行：python train_bearing.py。
2. 程序输出测试集准确率，并生成 bearing_model.h 头文件。
3. 在边缘设备工程中包含该头文件，调用 bearing_inference() 函数，传入16维归一化频域特征数组，返回值大于0.5表示故障。

判定标准：

• 训练输出准确率应高于0.85（由于随机初始化，每次结果略有波动，但应稳定在0.85~1.0之间）。
• 生成的 bearing_model.h 中权重数组非空且数值在[-128,127]范围内。
• 在PC上编写一个测试程序，将训练集第一个样本的16个特征值硬编码传入 bearing_inference()，输出概率与Python推理结果误差应小于0.05。