凯斯西储数据集(CWRU)十分类处理与训练代码(Pytorch)

Posted 2022-12-26 絮落潺流

数据预处理（十分类）
文件名：data_10

Fault location	Loads(hp)	Defect diameters (inches)	Class
Normal	0/1/2/3	0	0
Inner race	0/1/2/3	0.007 0.014 0.021	1 2 3
Ball	0/1/2/3	0.007 0.014 0.021	4 5 6
Outer race	0/1/2/3	0.007 0.014 0.021	7 8 9

import random
import numpy as np
import scipy.io as scio
from sklearn import preprocessing

def open_data(bath_path,key_num):
    #open_data('/Users/apple/Desktop/cwru/12k Drive End Bearing Fault Data/',105)
    path = bath_path + str(key_num) + ".mat"
    str1 =  "X" + "%03d"%key_num + "_DE_time"
    data = scio.loadmat(path)
    data = data[str1]
    return data


def deal_data(data,length,label):
    #line:num column:1025
    data = np.reshape(data,(-1))
    num = len(data)//length
    data = data[0:num*length]

    data = np.reshape(data,(num,length))

    min_max_scaler = preprocessing.MinMaxScaler()

    data = min_max_scaler.fit_transform(np.transpose(data,[1,0]))
    data = np.transpose(data,[1,0])
    label = np.ones((num,1))*label
    return np.column_stack((data,label)) 


def split_data(data,split_rate):
    length = len(data)
    num1 = int(length*split_rate[0])
    num2 = int(length*split_rate[1])

    index1 = random.sample(range(num1),num1)
    train = data[index1]
    data = np.delete(data,index1,axis=0)
    index2 = random.sample(range(num2),num2)
    valid = data[index2]
    test = np.delete(data,index2,axis=0)
    return train,valid,test


def load_data(num,length,hp,fault_diameter,split_rate):
    #num: number of sample in each data file
    #length: each sample
    #split_rate: train:valid:test
    bath_path1 = 'path of Normal Baseline Data'
    bath_path2 = 'path of 12k Drive End Bearing Fault Data/'
    data_list = []
    file_list = np.array([[105,118,130,106,119,131,107,120,132,108,121,133],  #0.007
                         [169,185,197,170,186,198,171,187,199,172,188,200],  #0.014
                         [209,222,234,210,223,235,211,224,236,212,225,237]])  #0.021
    label = 0
    #normal data
    for i in hp:
        normal_data = open_data(bath_path1,97+i)
        data = deal_data(normal_data,length,label = label)
        data_list.append(data)
        
    #abnormal data
    for i in fault_diameter:
        for j in hp:
            inner_num = file_list[int(i/0.007-1),3*j]
            ball_num = file_list[int(i/0.007-1),3*j+1]
            outer_num = file_list[int(i/0.007-1),3*j+2]
            

            inner_data = open_data(bath_path2,inner_num)
            inner_data = deal_data(inner_data,length,label + 1)
            data_list.append(inner_data)

            ball_data = open_data(bath_path2,ball_num)
            ball_data = deal_data(ball_data,length,label + 4)
            data_list.append(ball_data)

            outer_data = open_data(bath_path2,outer_num)
            outer_data = deal_data(outer_data,length,label + 7)
            data_list.append(outer_data)

        label = label + 1

    #keep each class same number of data
    num_list = []
    for i in data_list:
        num_list.append(len(i))
    min_num = min(num_list)

    if num > min_num:
        print("The number of each class overflow, the maximum number is：%d" %min_num)

    min_num = min(num,min_num)
    #Divide the train, validation, test sets and shuffle
    train = []
    valid = []
    test = []
    for data in data_list:
        data = data[0:min_num,:]
        a,b,c = split_data(data,split_rate)
        train.append(a)
        valid.append(b)
        test.append(c)

    train = np.reshape(train,(-1,length+1))
    train = train[random.sample(range(len(train)),len(train))]
    train_data = train[:,0:length]
    train_label = train[:,length]
    onehot_encoder = preprocessing.OneHotEncoder(sparse=False)
    train_label = train_label.reshape(len(train_label), 1)
    train_label = onehot_encoder.fit_transform(train_label)

    valid = np.reshape(valid,(-1,length+1))
    valid = valid[random.sample(range(len(valid)),len(valid))]
    valid_data = valid[:,0:length]
    valid_label = valid[:,length]
    valid_label = valid_label.reshape(len(valid_label), 1)
    valid_label = onehot_encoder.fit_transform(valid_label)

    test = np.reshape(test,(-1,length+1))
    test = test[random.sample(range(len(test)),len(test))]
    test_data = test[:,0:length]
    test_label = test[:,length]
    test_label = test_label.reshape(len(test_label), 1)
    test_label = onehot_encoder.fit_transform(test_label)


    return train_data,train_label,valid_data,valid_label,test_data,test_label
    

WDCNN训练与测试

from data_10 import load_data

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

import torch.utils.data
from torch.utils.data import DataLoader, TensorDataset
import torchvision
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self, in_channel=1, out_channel=10):
        super(TeacherNet, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv1d(in_channel, 16, kernel_size=64,stride=16,padding=24),  
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2,stride=2)
            )

        self.layer2 = nn.Sequential(
            nn.Conv1d(16, 32, kernel_size=3,padding=1), 
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2))  

        self.layer3 = nn.Sequential(
            nn.Conv1d(32, 64, kernel_size=3,padding=1),  
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )  

        self.layer4 = nn.Sequential(
            nn.Conv1d(64, 64, kernel_size=3,padding=1),  
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )  

        self.layer5 = nn.Sequential(
            nn.Conv1d(64, 64, kernel_size=3),  
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )  

        self.fc=nn.Sequential(
            nn.Linear(64, 100),
            nn.ReLU(inplace=True),
            nn.Linear(100, out_channel)
        )

    def forward(self, x):
        x = self.layer1(x) 
        x = self.layer2(x)  
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = x.view(x.size(0), -1)
        output = self.fc(x)
        return output

def train_Model(model,train_loader,optimizer,epoch):
    model.train()
    trained_samples = 0
    correct = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        
        optimizer.zero_grad()
        output = model(data)
        
        loss_fn = nn.MSELoss(reduce=True, size_average=True)
        loss = loss_fn(output.float(), target.float()) 
        loss.backward(loss.clone().detach())
        optimizer.step()
        
        trained_samples += len(data)
        
        print("\\rTrain epoch %d: %d/%d, " %
              (epoch, trained_samples, len(train_loader.dataset),), end='')
   
        pred = output.argmax(dim=1, keepdim=True)
        real = target.argmax(dim=1, keepdim=True)
        correct += pred.eq(real.view_as(pred)).sum().item()

    train_acc = correct / len(train_loader.dataset)
    print("Train acc: " , train_acc)


def test_Model(model,test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)  #logits
            print(output)
            loss_fn = nn.MSELoss(reduce=True, size_average=False)
            test_loss += loss_fn(output.float(), target.float()).item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
           # print((pred==4).sum())
            target = target.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    
    print('\\nTest: average loss: :.4f, accuracy: / (:.2f%)'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    
    return test_loss, correct / len(test_loader.dataset)

def main():
    epochs = 300
    batch_size = 32
    torch.manual_seed(0)
    
    train_dataset,train_label,_,_,test_dataset,test_label = load_data(num = 100,length = 1024,hp = [0,1,2,3],fault_diameter = [0.007,0.014,0.021],split_rate = [0.7,0.1,0.2])
    
    train_dataset = torch.tensor(train_dataset)
    train_label = torch.tensor(train_label)
    test_dataset = torch.tensor(test_dataset)
    test_label = torch.tensor(test_label)
    
    train_dataset = train_dataset.unsqueeze(1)

    test_dataset  = test_dataset.unsqueeze(1)
    train_dataset = train_dataset.to(torch.float32)
    test_dataset  = test_dataset.to(torch.float32)
    
    
    train_id = TensorDataset(train_dataset, train_label) 
    test_id  = TensorDataset(test_dataset, test_label)
    
    train_loader = DataLoader(dataset=train_id, batch_size=batch_size, shuffle=True)
    test_loader  = DataLoader(dataset=test_id,  batch_size=batch_size, shuffle=False)

    model = Net()
    optimizer = torch.optim.Adadelta(model.parameters())
    
    model_history = []

    for epoch in range(1, epochs + 1):
        train_Model(model, train_loader, optimizer, epoch)  
        loss, acc = test_Model(model, test_loader)
        model_history.append((loss, acc))

   # torch.save(model.state_dict(), "model.pt")
    return model, model_history

model, model_history = main()