最近这周也在想如何结合公司的业务做一个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 torchimport torch.nn as nnimport torchaudioimport numpy as npfrom sklearn.model_selection import train_test_splitfrom sklearn.preprocessing import LabelEncoderfrom 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 srcclass 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 和对应的标签列表 labelsaudio_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和DataLoadertrain_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_melsd_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.001device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果有GPU则使用GPU，否则使用CPUmodel.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 normalizationcorrect = 0total = 0with 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)