基于简单CNN的西楚数据集智能诊断

仅作为初学者入门的示范案例，其中，模型权重保存、加载、测试、可视化等代码略。

import

import torch
import numpy as np
import torch.nn as nn
import scipy.io as sio
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import maxabs_scale

数据集构建

class CWRU_DATA(Dataset):
    def __init__(self, data_num=100, data_len=1024):  # not overlapped
        self.data, self.label = self.prepare_data(num=data_num, length=data_len)  # (4, 100, 1024)

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

    def data_preprocess(self, x):
        x = x.astype(np.float32)  # shape: (N,)
        x = maxabs_scale(x)
        return x

    def prepare_data(self, num, length):
        NC = r"G:\dataset\casedata_12khz\normal\normal_0_X097.mat"
        IF = r"G:\dataset\casedata_12khz\inner\inner_007\inner007_0_105.mat"
        OF = r"G:\dataset\casedata_12khz\outer\outer_007\outer007_0.mat"
        BF = r"G:\dataset\casedata_12khz\ball\ball_007\ball007_0.mat"
        paths = [NC, IF, OF, BF]
        data = np.zeros([4, num, length], dtype=np.float32)
        for i in range(len(data)):
            data[i] = self.load_mat_data(paths[i], item_name='DE_time',
                                         num_points=num * length).reshape([num, length])
        label = np.arange(len(paths)).reshape([len(paths), 1])
        label = np.tile(label, (1, num))  # (4, num)

        return data.reshape(-1, length), label.reshape(-1)

    def load_mat_data(self, path, item_name='DE_time', num_points=1024 * 100):
        # get data from .mat file
        mat_data = sio.loadmat(path)  # a dictionary
        name_new = []
        for n in list(mat_data.keys()):
            if item_name in n:
                name_new.append(n)
        data = mat_data[name_new[0]]
        # TODO: you can save these data in another format, such as .npy, .csv, .txt
        return data[:num_points]

    def __getitem__(self, item):
        data, label = self.data[item], self.label[item]  # (1024, ), (1, )
        data = self.data_preprocess(data)
        return data, label

Dataloader创建

def build_dataloader(dataset, shuffle, batch_size=32):
    return DataLoader(dataset, batch_size, shuffle=shuffle)

Model 构建

def conv_block(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm1d(out_channels),
        nn.ReLU(),
        nn.MaxPool1d(kernel_size=2),
    )


class SimpleCNN(nn.Module):
    def __init__(self, in_chn=1, hidden_chn=64, block_num=4, sample_len=1024, classes=4):
        super().__init__()
        conv1 = conv_block(in_chn, hidden_chn)
        conv1_more = [conv_block(hidden_chn, hidden_chn) for _ in range(block_num - 1)]
        self.features = nn.Sequential(conv1, *conv1_more)  # (None, 64, 1024/2^4)
        flatten_size = int(sample_len / (2 ** block_num) * hidden_chn)
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(flatten_size, classes))

    def forward(self, x):
        feat = self.features(x)
        out = self.classifier(feat)
        return out

训练

if __name__ == "__main__":

    dataset = CWRU_DATA()
    dataloader = build_dataloader(dataset, shuffle=True, batch_size=32)

    model = SimpleCNN()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss(reduction='mean')
    # This criterion combines torch.nn.LogSoftmax and torch.nn.NLLLoss in one single class

    for i, batch in enumerate(dataloader):
        data, label = batch
        data, label = data.float().unsqueeze(1), label.long()
        predict_label = model(data)
        loss = criterion(predict_label, label)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        print(loss.detach().item())