最近这周也在想如何结合公司的业务做一个embedding(高维向量转换模型), 为风电业务接入RAG系统做准备. transformer可以自己定义 or 采用科大讯飞的，以下只是一种实现方式, 仅供参考
概述
基于深度学习的电机振动声音分类系统，该系统使用Transformer架构对电机振动音频进行分析，以识别不同的电机状态（如正常运行、轴承故障、不平衡等）。系统结合了PyTorch深度学习框架和Scikit-learn机器学习库的功能，实现了从数据预处理到模型训练、评估和故障诊断的完整流程。
系统架构
系统由以下几个主要组件构成：
1. 数据加载与预处理
自定义数据集类
系统使用自定义的VibrationDataset类来加载和处理电机振动音频数据。继承自PyTorch的Dataset类，实现了必要的__len__和__getitem__方法。
主要功能：
● 加载音频文件及其对应的标签
● 应用特征转换（如梅尔频谱图提取）
● 返回处理后的特征和标签对
标签编码
系统使用Scikit-learn的LabelEncoder将文本标签（如"normal"、"bearing_fault"等）转换为数值形式，便于模型处理。这与Scikit-learn中的标准预处理流程一致，如文档中所述，预处理是机器学习流程中的重要步骤。
数据集划分
使用Scikit-learn的train_test_split函数将数据集划分为训练集和测试集，这是模型评估的标准做法。该函数确保数据被随机且均匀地分配到训练集和测试集中，同时通过设置random_state参数保证结果的可重复性。
2. 特征提取
梅尔频谱图提取
系统使用自定义的MelSpectrogram类（基于PyTorch的nn.Module）从原始音频波形中提取梅尔频谱图特征。梅尔频谱图是音频处理中常用的特征表示方法，能够更好地模拟人耳对声音的感知。
主要处理步骤：
1. 使用短时傅里叶变换(STFT)将时域信号转换为频域
2. 应用梅尔滤波器组
3. 对结果取对数，增强低能量部分的特征
4. 转置结果，使时间成为序列维度
3. 模型架构
Transformer编码器层
系统实现了自定义的TransformerEncoderLayer类，作为Transformer编码器的基本构建块。该层包含：
● 多头自注意力机制
● 前馈神经网络
● 残差连接
● 层归一化
● Dropout正则化
Transformer分类器
TransformerClassifier类是整个模型的核心，它将多个Transformer编码器层堆叠起来，并添加了：
● 输入嵌入层：将特征维度映射到模型维度
● 序列池化：对时间维度取平均，得到固定长度的表示
● 分类头：最终的线性分类层
这种架构设计使模型能够有效捕获音频特征序列中的时间依赖关系，适合处理电机振动这类时序数据。
4. 模型训练与评估
训练流程
训练过程遵循标准的深度学习训练流程：
-
1. 设置模型为训练模式
-
2. 迭代训练数据加载器
-
3. 前向传播计算预测
-
4. 计算损失（使用交叉熵损失函数）
-
5. 反向传播计算梯度
-
6. 更新模型参数
-
7. 记录和打印训练进度
评估方法
模型评估采用准确率（Accuracy）作为主要指标，与Scikit-learn中的评估方法类似。评估过程包括：
-
1. 设置模型为评估模式
-
2. 禁用梯度计算
-
3. 对测试集数据进行预测
-
4. 计算预测准确率
这种评估方法符合Scikit-learn中的模型评估最佳实践，
5. 故障诊断应用
系统提供了predict_and_suggest函数，用于对单个音频样本进行预测并给出诊断建议。该函数：
-
1. 加载音频文件
-
2. 提取特征
-
3. 使用训练好的模型进行预测
-
4. 将预测结果（数字标签）转换回原始文本标签
-
5. 根据预测的故障类型提供相应的诊断建议
这种应用方式展示了机器学习模型在实际工业场景中的应用价值。
实现细节
数据预处理参数
系统使用以下参数进行音频特征提取：
● 采样率：16000 Hz
● FFT窗口大小：400
● FFT窗口步长：160
● 梅尔滤波器组数量：64
模型超参数
Transformer模型使用以下超参数：
● 模型维度：128
● 注意力头数：4
● 编码器层数：2
● Dropout概率：0.1
● 批量大小：32
● 学习率：0.001
● 训练轮数：10
使用指南
数据准备
使用系统前，需要准备：
1. 电机振动音频文件
2. 对应的标签（如"normal"、"bearing_fault"等）
模型训练
按照代码中的流程，可以完成模型的训练和评估：
1. 准备音频文件路径和标签
2. 创建数据集和数据加载器
3. 初始化模型
4. 训练模型
5. 评估模型性能
故障诊断
使用训练好的模型进行故障诊断：
test_audio_path = "path/to/your/test_audio.wav"
predict_and_suggest(test_audio_path, model, mel_transform, label_encoder, device)
完整代码
import torch
import torch.nn as nn
import torchaudio
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from torch.utils.data import Dataset, DataLoader
# 1. 数据加载与预处理
class VibrationDataset(Dataset):
"""
自定义数据集类，用于加载电机振动音频数据及其对应的标签。
继承自torch.utils.data.Dataset，需要实现__len__和__getitem__方法。
"""
def __init__(self, audio_paths, labels, transform=None):
"""
初始化VibrationDataset。
Args:
audio_paths (list): 音频文件路径的列表。
labels (list): 与音频文件对应的标签列表。
transform (callable, optional): 对音频数据进行预处理的转换函数。默认为None。
"""
self.audio_paths = audio_paths
self.labels = labels
self.transform = transform
def __len__(self):
"""
返回数据集的总大小。
"""
return len(self.audio_paths)
def __getitem__(self, idx):
"""
根据给定的索引获取一个数据样本。
Args:
idx (int): 数据样本的索引。
Returns:
tuple: 包含特征 (音频数据的转换结果) 和标签的元组。
"""
audio_path = self.audio_paths[idx]
label = self.labels[idx]
waveform, sample_rate = torchaudio.load(audio_path) # 使用torchaudio加载音频文件，返回波形和采样率
if self.transform:
features = self.transform(waveform, sample_rate) # 如果提供了转换函数，则对波形进行转换
return features, label
# 特征提取 (例如 Mel-Spectrogram)
class MelSpectrogram(nn.Module):
"""
定义一个提取梅尔频谱特征的PyTorch模块。
继承自torch.nn.Module。
"""
def __init__(self, sample_rate, n_fft, hop_length, n_mels):
"""
初始化MelSpectrogram模块。
Args:
sample_rate (int): 音频的采样率。
n_fft (int): FFT窗口的大小。
hop_length (int): 相邻FFT窗口之间的步长。
n_mels (int): 梅尔滤波器组的数量。
"""
super().__init__()
self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
sample_rate=sample_rate,
n_fft=n_fft,
hop_length=hop_length,
n_mels=n_mels
)
def forward(self, waveform, sample_rate):
"""
定义前向传播过程，将原始波形转换为梅尔频谱。
Args:
waveform (torch.Tensor): 原始音频波形，形状为 (num_channels, num_samples)。
sample_rate (int): 音频的采样率 (这里可能不直接使用，因为在初始化时已指定)。
Returns:
torch.Tensor: 梅尔频谱，形状为 (time, n_mels)。
"""
mel_spec = self.mel_spectrogram(waveform) # 计算梅尔频谱
# 可以进行一些额外的处理，例如对数变换以增强低能量部分的特征
mel_spec = torch.log(mel_spec + 1e-6) # 加一个小常数以避免log(0)
return mel_spec.transpose(0, 1) # 将形状从 (n_mels, time) 转置为 (time, n_mels)，使时间成为序列维度
# 2. Transformer模型定义
class TransformerEncoderLayer(nn.Module):
"""
Transformer编码器中的一个单独的层。
包含多头自注意力机制和前馈神经网络，并带有残差连接和层归一化。
"""
def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
"""
初始化TransformerEncoderLayer。
Args:
d_model (int): 输入特征的维度 (也即Transformer模型中使用的维度)。
nhead (int): 多头自注意力机制中注意力头的数量。
dim_feedforward (int): 前馈神经网络中间层的维度。
dropout (float): dropout的概率。
"""
super().__init__()
self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) # 多头自注意力
self.linear1 = nn.Linear(d_model, dim_feedforward) # 第一个线性层
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(dim_feedforward, d_model) # 第二个线性层
self.norm1 = nn.LayerNorm(d_model) # 第一个层归一化
self.norm2 = nn.LayerNorm(d_model) # 第二个层归一化
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.activation = nn.ReLU() # 激活函数
def forward(self, src, src_mask=None, src_key_padding_mask=None):
"""
定义TransformerEncoderLayer的前向传播过程。
Args:
src (torch.Tensor): 输入序列，形状为 (batch_size, seq_len, d_model)。
src_mask (torch.Tensor, optional): 注意力掩码，用于屏蔽某些位置。默认为None。
src_key_padding_mask (torch.Tensor, optional): 用于屏蔽padding的key。默认为None。
Returns:
torch.Tensor: 经过编码器层处理后的输出，形状为 (batch_size, seq_len, d_model)。
"""
attn_output, _ = self.self_attn(src, src, src, attn_mask=src_mask,
key_padding_mask=src_key_padding_mask, is_causal=False)
# 残差连接 + 层归一化 + dropout
src = self.norm1(src + self.dropout1(attn_output))
ff_output = self.linear2(self.dropout(self.activation(self.linear1(src))))
# 残差连接 + 层归一化 + dropout
src = self.norm2(src + self.dropout2(ff_output))
return src
class TransformerClassifier(nn.Module):
"""
基于Transformer编码器的分类模型。
将输入的特征序列通过Transformer编码器处理后，进行分类。
"""
def __init__(self, input_dim, d_model, nhead, num_layers, num_classes, dim_feedforward=2048, dropout=0.1):
"""
初始化TransformerClassifier。
Args:
input_dim (int): 输入特征的维度 (例如，梅尔频谱的n_mels)。
d_model (int): Transformer模型中使用的维度。
nhead (int): 多头自注意力机制中注意力头的数量。
num_layers (int): Transformer编码器层的数量。
num_classes (int): 分类的类别数量。
dim_feedforward (int): 前馈神经网络中间层的维度。
dropout (float): dropout的概率。
"""
super().__init__()
self.embedding = nn.Linear(input_dim, d_model) # 将输入特征维度映射到Transformer的维度
self.transformer_encoder = nn.ModuleList([
TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout) for _ in range(num_layers)
]) # 堆叠多个Transformer编码器层
self.linear = nn.Linear(d_model, num_classes) # 最终的线性分类层
self.dropout = nn.Dropout(dropout)
self.norm = nn.LayerNorm(d_model) # 最后的层归一化
def forward(self, src, src_mask=None, src_key_padding_mask=None):
"""
定义TransformerClassifier的前向传播过程。
Args:
src (torch.Tensor): 输入特征序列，形状为 (batch_size, seq_len, input_dim)。
src_mask (torch.Tensor, optional): 注意力掩码。默认为None。
src_key_padding_mask (torch.Tensor, optional): 用于屏蔽padding的key。默认为None。
Returns:
torch.Tensor: 分类模型的输出，形状为 (batch_size, num_classes)。
"""
src = self.embedding(src) # (batch_size, seq_len, d_model) - 将输入维度嵌入到模型维度
for layer in self.transformer_encoder:
src = layer(src, src_mask, src_key_padding_mask) # 通过每个Transformer编码器层
# 可以选择对序列的哪个部分进行分类，例如最后一个时间步的输出，或者对所有时间步的输出进行池化
# 这里简单地取序列的平均作为整个序列的表示
pooled_output = torch.mean(src, dim=1) # (batch_size, d_model) - 对时间序列维度求平均
output = self.linear(self.dropout(self.norm(pooled_output))) # 通过线性层进行分类
return output
# 3. 数据准备
# 假设你已经有了音频文件路径列表 audio_paths 和对应的标签列表 labels
audio_paths = [...] # 替换为你的音频文件路径列表
labels = [...] # 替换为你的音频文件对应的标签列表
# 对标签进行编码，将文本标签转换为数字
label_encoder = LabelEncoder()
encoded_labels = label_encoder.fit_transform(labels)
num_classes = len(label_encoder.classes_) # 获取类别数量
# 划分训练集和测试集
train_paths, test_paths, train_labels, test_labels = train_test_split(
audio_paths, encoded_labels, test_size=0.2, random_state=42 # 80%训练，20%测试，设置随机种子以保证可重复性
)
# 定义特征提取参数
sample_rate = 16000 # 音频采样率
n_fft = 400 # FFT窗口大小
hop_length = 160 # FFT窗口步长
n_mels = 64 # 梅尔滤波器组数量
mel_transform = MelSpectrogram(sample_rate, n_fft, hop_length, n_mels) # 实例化梅尔频谱转换
# 创建Dataset和DataLoader
train_dataset = VibrationDataset(train_paths, train_labels, transform=mel_transform)
test_dataset = VibrationDataset(test_paths, test_labels, transform=mel_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) # 训练集DataLoader，shuffle=True表示每个epoch都打乱数据
test_loader = DataLoader(test_dataset, batch_size=32) # 测试集DataLoader
# 4. 模型初始化
input_dim = n_mels # 输入特征维度是梅尔频谱的n_mels
d_model = 128 # Transformer模型内部的维度
nhead = 4 # 多头注意力的头数
num_layers = 2 # Transformer编码器层的数量
dropout = 0.1 # dropout概率
model = TransformerClassifier(input_dim, d_model, nhead, num_layers, num_classes, dropout=dropout) # 实例化Transformer分类器
# 5. 训练模型
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数，适用于多分类问题
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam优化器，学习率为0.001
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果有GPU则使用GPU，否则使用CPU
model.to(device) # 将模型移动到指定的设备
num_epochs = 10 # 训练的轮数
for epoch in range(num_epochs):
model.train() # 设置模型为训练模式
for inputs, labels in train_loader:
inputs = inputs.to(device) # 将输入数据移动到指定设备
labels = labels.to(device) # 将标签移动到指定设备
optimizer.zero_grad() # 清空之前的梯度
outputs = model(inputs) # 前向传播
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新模型参数
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
# 6. 模型评估
model.eval() # 设置模型为评估模式，禁用dropout和batch normalization
correct = 0
total = 0
with torch.no_grad(): # 在评估模式下禁用梯度计算
for inputs, labels in test_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1) # 获取每个样本预测概率最大的类别
total += labels.size(0) # 统计总样本数
correct += (predicted == labels).sum().item() # 统计预测正确的样本数
print(f"Accuracy of the model on the test set: {100 * correct / total:.2f}%")
# 7. 诊断建议生成 (简化示例)
def predict_and_suggest(audio_path, model, transform, label_encoder, device):
"""
对单个音频文件进行预测并给出简单的诊断建议。
Args:
audio_path (str): 要预测的音频文件路径。
model (nn.Module): 训练好的模型。
transform (callable): 用于提取特征的转换函数。
label_encoder (LabelEncoder): 用于将数字标签转换回文本标签。
device (torch.device): 模型所在的设备。
"""
model.eval() # 设置模型为评估模式
waveform, sample_rate = torchaudio.load(audio_path)
features = transform(waveform, sample_rate).unsqueeze(0).to(device) # 提取特征并添加batch维度，移动到指定设备
with torch.no_grad(): # 禁用梯度计算
output = model(features)
_, predicted_idx = torch.max(output, 1) # 获取预测的类别索引
predicted_label = label_encoder.inverse_transform([predicted_idx.item()])[0] # 将索引转换回文本标签
print(f"Predicted state: {predicted_label}")
# 根据预测结果给出诊断建议 (需要更完善的故障知识库)
if predicted_label == "normal":
print("建议: 电机运行正常。")
elif predicted_label == "bearing_fault":
print("建议: 检测轴承是否存在磨损或损坏，建议进行润滑或更换。")
elif predicted_label == "imbalance":
print("建议: 检查电机转子是否平衡，可能需要进行动平衡校正。")
# ... 可以添加更多故障类型的诊断建议
# 使用训练好的模型进行预测和建议 (你需要替换为实际的音频文件路径)
test_audio_path = "path/to/your/test_audio.wav"
predict_and_suggest(test_audio_path, model, mel_transform, label_encoder, device)
打包模型
import os
import json
import torch
from huggingface_hub import HfApi, HfFolder
# 创建保存目录
model_dir = "motor_vibration_classifier"
os.makedirs(model_dir, exist_ok=True)
# 1. 保存模型
torch.save(model.state_dict(), os.path.join(model_dir, "model.pt"))
# 2. 保存标签编码器
import pickle
with open(os.path.join(model_dir, "label_encoder.pkl"), "wb") as f:
pickle.dump(label_encoder, f)
# 3. 保存特征提取器参数
mel_config = {
"sample_rate": sample_rate,
"n_fft": n_fft,
"hop_length": hop_length,
"n_mels": n_mels
}
with open(os.path.join(model_dir, "mel_config.json"), "w") as f:
json.dump(mel_config, f)
# 4. 保存模型配置
model_config = {
"input_dim": input_dim,
"d_model": d_model,
"nhead": nhead,
"num_layers": num_layers,
"num_classes": num_classes,
"dim_feedforward": 2048,
"dropout": dropout
}
with open(os.path.join(model_dir, "model_config.json"), "w") as f:
json.dump(model_config, f)
