基于卷积神经网络CNN实现旋转机械故障诊断——实现数据集：CWRU西储大学轴承，JN江南大学轴承，东南大学齿轮，HUST华科轴承注：Pytorch和tensorflow两种框架下实现，提供一种实现流程

计算机c9硕士算法工程师

1510人浏览 · 2024-10-26 20:45:10

计算机c9硕士算法工程师 · 2024-10-26 20:45:10 发布

使用卷积神经网络（CNN）来实现旋转机械的故障诊断。具体来说，该项目将使用四个不同的数据集：CWRU（西储大学）轴承数据集、江南大学（JN）轴承数据集、东南大学齿轮数据集和华中科技大学（HUST）轴承数据集。我们将使用PyTorch和TensorFlow两种深度学习框架来实现这一目标，以便读者可以根据自己的喜好选择适合的框架。

数据集

CWRU（西储大学）轴承数据集
- 描述：包含多种不同状态下的滚动轴承振动信号。
- 用途：滚动轴承故障诊断。
- 大小：约1GB。
江南大学（JN）轴承数据集
- 描述：包含多种不同工况下的滚动轴承振动信号。
- 用途：滚动轴承故障诊断。
- 大小：约1GB。
东南大学齿轮数据集
- 描述：包含多种不同状态下的齿轮振动信号。
- 用途：齿轮故障诊断。
- 大小：约1GB。
华中科技大学（HUST）轴承数据集
- 描述：包含多种不同状态下的滚动轴承振动信号。
- 用途：滚动轴承故障诊断。
- 大小：约1GB。

实现流程

数据预处理：将数据集转换为适合CNN输入的格式。
构建CNN模型：使用PyTorch或TensorFlow构建卷积神经网络模型。
训练模型：使用训练数据集训练CNN模型。
评估模型：使用测试数据集评估模型的性能。
结果分析：分析模型的准确率、损失函数等性能指标。

项目目录结构

RotaryMachineFaultDiagnosis/
├── data/
│   ├── CWRU_Bearing_Dataset.zip
│   ├── JN_Bearing_Dataset.zip
│   ├── SEU_Gear_Dataset.zip
│   ├── HUST_Bearing_Dataset.zip
├── src/
│   ├── pytorch/
│   │   ├── main.py
│   │   ├── model.py
│   │   ├── dataset.py
│   ├── tensorflow/
│   │   ├── main.py
│   │   ├── model.py
│   │   ├── dataset.py
└── README.md  # 项目说明

示例代码

PyTorch实现

数据集定义 (`src/pytorch/dataset.py`)

import torch
from torch.utils.data import Dataset
import os
import numpy as np

class BearingDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.file_list = os.listdir(data_dir)
        self.transform = transform

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

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        file_path = os.path.join(self.data_dir, self.file_list[idx])
        signal = np.load(file_path)
        label = int(self.file_list[idx].split('_')[1])  # 假设标签在文件名中
        if self.transform:
            signal = self.transform(signal)
        return signal, label

CNN模型定义 (`src/pytorch/model.py`)

import torch
import torch.nn as nn

class CNNClassifier(nn.Module):
    def __init__(self):
        super(CNNClassifier, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool1d(kernel_size=2)
        self.conv2 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool1d(kernel_size=2)
        self.fc1 = nn.Linear(32 * 128, 128)
        self.fc2 = nn.Linear(128, 10)  # 假设10类故障

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        x = x.view(-1, 32 * 128)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

主程序 (`src/pytorch/main.py`)

import torch
from torch.utils.data import DataLoader
from torch import optim
import torch.nn.functional as F
from dataset import BearingDataset
from model import CNNClassifier

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

dataset = BearingDataset('data/CWRU_Bearing_Dataset')
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

model = CNNClassifier().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 20
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.float().unsqueeze(1).to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/(batch_idx+1):.4f}')

TensorFlow实现

数据集定义 (`src/tensorflow/dataset.py`)

import tensorflow as tf
import numpy as np
import os

def parse_npz(file_path):
    signal = np.load(file_path)
    label = int(file_path.split('_')[1])
    return signal, label

def create_dataset(data_dir, batch_size=32):
    file_paths = [os.path.join(data_dir, fname) for fname in os.listdir(data_dir)]
    dataset = tf.data.Dataset.from_tensor_slices(file_paths)
    dataset = dataset.map(lambda x: tf.numpy_function(parse_npz, [x], Tout=(tf.float32, tf.int32)))
    dataset = dataset.batch(batch_size)
    return dataset

CNN模型定义 (`src/tensorflow/model.py`)

import tensorflow as tf
from tensorflow.keras import layers

def build_cnn_classifier():
    model = tf.keras.Sequential([
        layers.Input(shape=(1024, 1)),
        layers.Conv1D(filters=16, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling1D(pool_size=2),
        layers.Conv1D(filters=32, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling1D(pool_size=2),
        layers.Flatten(),
        layers.Dense(128, activation='relu'),
        layers.Dense(10)  # 假设10类故障
    ])
    return model

主程序 (`src/tensorflow/main.py`)

import tensorflow as tf
from dataset import create_dataset
from model import build_cnn_classifier

dataset = create_dataset('data/CWRU_Bearing_Dataset')
model = build_cnn_classifier()
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

num_epochs = 20
history = model.fit(dataset, epochs=num_epochs)

项目运行

确保安装了必要的依赖库：

pip install torch torchvision tensorflow

然后分别运行PyTorch和TensorFlow的主程序：

python src/pytorch/main.py
python src/tensorflow/main.py

学习资源

项目中的代码包含了详细的注释，帮助初学者理解各个部分的功能和作用。同时，提供的数据集可以让用户快速上手，了解如何使用卷积神经网络进行旋转机械故障诊断。

总结

这个旋转机械故障诊断项目是一个完整的解决方案，它不仅包含了多种数据集，还包括了使用PyTorch和TensorFlow两种深度学习框架实现的代码。通过本项目，你可以深入学习如何使用CNN进行旋转机械故障诊断，并将其应用于实际的故障诊断中。对于初学者来说，这是一个很好的学习平台，可以深入了解卷积神经网络及其在故障诊断中的应用。