基于双路神经网络的滚动轴承故障诊断融合了原始振动信号 和 二维信号时频图像 的多输入(多通道)故障诊断方法单路和双路都可时频图像算法可选小波变换，短时傅里叶变换，马尔可夫变迁场，格拉姆角场故障诊断

基于双路神经网络的滚动轴承故障诊断

融合了原始振动信号 和 二维信号时频图像 的多输入(多通道)故障诊断方法

单路和双路都可

时频图像算法可选小波变换，短时傅里叶变换，马尔可夫变迁场，格拉姆角场，S变换，递归图，灰度图等基于双路神经网络的滚动轴承故障诊断

融合了原始振动信号 和 二维信号时频图像 的多输入(多通道)故障诊断方法

单路和双路都可

时频图像算法可选小波变换，短时傅里叶变换，马尔可夫变迁场，格拉姆角场，S变换，递归图，灰度图等

数据集可选凯斯西储大学，东南大学，江南大学等等

Pytorch和Tensorflow框架做的。

以下可视化结果适用小波和短时傅里叶变换，其它时频图像算法效果并没有这么好，我提供框架，具体任需自己调优

Pytorch和Tensorflow框架制作。

基于双路神经网络的滚动轴承故障诊断
类别：深度学习、故障诊断、信号处理、多通道输入
用途：该项目旨在通过融合原始振动信号和二维时频图像来提高滚动轴承故障诊断的准确性。采用双路神经网络架构，支持单路和双路输入，并提供了多种时频图像算法（如小波变换、短时傅里叶变换等）。项目使用PyTorch和TensorFlow框架实现，并提供了一个完整的解决方案，包括数据集、模型训练和预测。

项目特点

1. 多通道输入：支持原始振动信号和二维时频图像的多通道输入。
2. 多种时频图像算法：提供多种时频图像生成方法，如小波变换、短时傅里叶变换、马尔可夫变迁场、格拉姆角场、S变换、递归图、灰度图等。
3. 双路神经网络架构：采用双路神经网络，分别处理原始振动信号和时频图像，然后融合特征进行故障诊断。
4. 灵活性：支持单路和双路输入，用户可以根据需要选择合适的输入方式。
5. 可视化结果：提供可视化工具，展示不同算法生成的时频图像及其在故障诊断中的效果。
6. 预训练模型：包含预训练模型，可以直接用于故障诊断。
7. 详细注释：代码中包含详细的注释，适合初学者学习和理解。

项目结构

bearing_fault_diagnosis/
├── data/  # 数据文件夹
│   ├── raw_vibration/  # 原始振动信号
│   ├── time_frequency_images/  # 时频图像
├── models/  # 模型文件夹
│   ├── single_channel_model.pth  # 单通道模型
│   ├── dual_channel_model.pth  # 双通道模型
├── train.py  # 训练脚本
├── predict.py  # 预测脚本
├── visualize.py  # 可视化脚本
├── requirements.txt  # 环境依赖列表
└── README.md  # 项目说明文档

• data/：存放数据集，包括原始振动信号和时频图像。
• models/：存放训练好的单通道和双通道模型。
• train.py：用于训练神经网络模型。
• predict.py：用于预测新的振动信号。
• visualize.py：用于生成和可视化时频图像。
• requirements.txt：环境依赖列表。
• README.md：项目说明文档，包含项目介绍、运行步骤和注意事项。

详细步骤

1. 安装依赖

确保您已经安装了Python 3.8及以上版本，并且有PyTorch 1.10.0或更高版本以及TensorFlow 2.4.0或更高版本。然后在项目根目录下运行以下命令安装依赖：

pip install -r requirements.txt

2. 数据集准备

将提供的数据集放置在data/目录下，确保目录结构如下：

data/
├── raw_vibration/  # 存放原始振动信号
├── time_frequency_images/  # 存放对应的时频图像

3. 生成时频图像

使用visualize.py脚本来生成时频图像。该脚本会读取原始振动信号，并使用指定的时频图像算法生成图像。

# visualize.py
import os
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as signal
from pyts.image import RecurrencePlot, GramianAngularField, MarkovTransitionField

def load_vibration_signal(file_path):
    # 加载原始振动信号
    return np.loadtxt(file_path)

def generate_time_frequency_image(signal, method='wavelet'):
    if method == 'wavelet':
        cwtmatr, freqs = signal.cwt(signal, signal.ricker, widths=np.arange(1, 10))
        return cwtmatr
    elif method == 'stft':
        f, t, Zxx = signal.stft(signal, fs=1000, nperseg=100)
        return np.abs(Zxx)
    elif method == 'recurrence':
        rp = RecurrencePlot()
        return rp.fit_transform([signal])[0]
    elif method == 'gramian':
        gaf = GramianAngularField()
        return gaf.fit_transform([signal])[0]
    elif method == 'markov':
        mtf = MarkovTransitionField()
        return mtf.fit_transform([signal])[0]
    else:
        raise ValueError("Unsupported method")

def save_image(image, file_path):
    plt.imshow(image, cmap='viridis')
    plt.axis('off')
    plt.savefig(file_path, bbox_inches='tight', pad_inches=0)
    plt.close()

def main():
    vibration_dir = 'data/raw_vibration'
    image_dir = 'data/time_frequency_images'
    
    for filename in os.listdir(vibration_dir):
        if filename.endswith(".txt"):
            file_path = os.path.join(vibration_dir, filename)
            signal = load_vibration_signal(file_path)
            
            # 生成并保存不同的时频图像
            for method in ['wavelet', 'stft', 'recurrence', 'gramian', 'markov']:
                image = generate_time_frequency_image(signal, method=method)
                image_file_path = os.path.join(image_dir, f"{os.path.splitext(filename)[0]}_{method}.png")
                save_image(image, image_file_path)
                print(f"Saved {image_file_path}")

if __name__ == '__main__':
    main()

4. 训练神经网络模型

使用train.py脚本来训练单通道和双通道神经网络模型。该脚本会读取数据集中的原始振动信号和时频图像，并开始训练过程。训练完成后，模型会被保存到models/目录下。

# train.py
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import numpy as np

class BearingDataset(Dataset):
    def __init__(self, vibration_dir, image_dir, transform=None):
        self.vibration_dir = vibration_dir
        self.image_dir = image_dir
        self.transform = transform
        self.files = os.listdir(vibration_dir)

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        vibration_file = os.path.join(self.vibration_dir, self.files[idx])
        image_file = os.path.join(self.image_dir, f"{os.path.splitext(self.files[idx])[0]}_wavelet.png")  # 以小波变换为例
        label = int(self.files[idx].split('_')[0])  # 假设文件名前缀为标签
        
        vibration = np.loadtxt(vibration_file).astype(np.float32)
        image = Image.open(image_file).convert('L')
        
        if self.transform:
            image = self.transform(image)
        
        return vibration, image, label

class SingleChannelModel(nn.Module):
    def __init__(self):
        super(SingleChannelModel, self).__init__()
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)  # 假设有10个类别

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class DualChannelModel(nn.Module):
    def __init__(self):
        super(DualChannelModel, self).__init__()
        self.fc1 = nn.Linear(1024, 512)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc2 = nn.Linear(512 + 32 * 64 * 64, 256)
        self.fc3 = nn.Linear(256, 10)  # 假设有10个类别

    def forward(self, vibration, image):
        vibration = torch.relu(self.fc1(vibration))
        
        image = self.pool(torch.relu(self.conv1(image)))
        image = image.view(image.size(0), -1)
        
        combined = torch.cat((vibration, image), dim=1)
        combined = torch.relu(self.fc2(combined))
        combined = self.fc3(combined)
        return combined

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 数据增强和预处理
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5])
    ])

    # 加载数据集
    dataset = BearingDataset('data/raw_vibration', 'data/time_frequency_images', transform=transform)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

    # 初始化模型
    single_channel_model = SingleChannelModel().to(device)
    dual_channel_model = DualChannelModel().to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer_single = optim.Adam(single_channel_model.parameters(), lr=0.001)
    optimizer_dual = optim.Adam(dual_channel_model.parameters(), lr=0.001)

    # 训练循环
    num_epochs = 100
    for epoch in range(num_epochs):
        for i, (vibration, image, label) in enumerate(dataloader):
            vibration, image, label = vibration.to(device), image.to(device), label.to(device)

            # 单通道模型
            output_single = single_channel_model(vibration)
            loss_single = criterion(output_single, label)
            optimizer_single.zero_grad()
            loss_single.backward()
            optimizer_single.step()

            # 双通道模型
            output_dual = dual_channel_model(vibration, image)
            loss_dual = criterion(output_dual, label)
            optimizer_dual.zero_grad()
            loss_dual.backward()
            optimizer_dual.step()

            if (i+1) % 10 == 0:
                print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(dataloader)}], Loss Single: {loss_single.item():.4f}, Loss Dual: {loss_dual.item():.4f}")

    # 保存模型
    torch.save(single_channel_model.state_dict(), 'models/single_channel_model.pth')
    torch.save(dual_channel_model.state_dict(), 'models/dual_channel_model.pth')
    print("Models saved.")

if __name__ == '__main__':
    main()

5. 使用预训练模型进行预测

等等