create three layer of CNN

Number-Of-CNN-Layers/_main_pytorch_threeLayerCNN.py

@@ -0,0 +1,233 @@
+import os
+import datetime
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchinfo import summary
+from torch.utils.tensorboard import SummaryWriter
+from sklearn.metrics import accuracy_score
+from tqdm import tqdm
+
+from utils import load_data, plot_history_torch, plot_heat_map
+
+# project root path
+project_path = "./"
+# define log directory
+# must be a subdirectory of the directory specified when starting the web application
+# it is recommended to use the date time as the subdirectory name
+log_dir = project_path + "logs/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+model_path = project_path + "ecg_model.pt"
+
+# the device to use
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print("Using {} device".format(device))
+
+
+# define the dataset class
+class ECGDataset(Dataset):
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+    def __getitem__(self, index):
+        x = torch.tensor(self.x[index], dtype=torch.float32)
+        y = torch.tensor(self.y[index], dtype=torch.long)
+        return x, y
+
+    def __len__(self):
+        return len(self.x)
+
+
+# build the CNN model
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # the first convolution layer, 4 21x1 convolution kernels, output shape (batch_size, 4, 300)
+        self.conv1 = nn.Conv1d(in_channels=1, out_channels=4, kernel_size=21, stride=1, padding='same')
+        # the first pooling layer, max pooling, pooling size=3 , stride=2, output shape (batch_size, 4, 150)
+        self.pool1 = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+        # the second convolution layer, 16 23x1 convolution kernels, output shape (batch_size, 16, 150)
+        self.conv2 = nn.Conv1d(in_channels=4, out_channels=16, kernel_size=23, stride=1, padding='same')
+        # the second pooling layer, max pooling, pooling size=3, stride=2, output shape (batch_size, 16, 75)
+        self.pool2 = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+        # the third convolution layer, 32 25x1 convolution kernels, output shape (batch_size, 32, 75)
+        self.conv3 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=25, stride=1, padding='same')
+        # the third pooling layer, average pooling, pooling size=3, stride=2, output shape (batch_size, 32, 38)
+        self.pool3 = nn.AvgPool1d(kernel_size=3, stride=2, padding=1)
+        # the fourth convolution layer, 64 27x1 convolution kernels, output shape (batch_size, 64, 38)
+        self.conv4 = nn.Conv1d(in_channels=32, out_channels=64, kernel_size=27, stride=1, padding='same')
+        # flatten layer, for the next fully connected layer, output shape (batch_size, 38*64)
+        self.flatten = nn.Flatten()
+        # fully connected layer, 128 nodes, output shape (batch_size, 128)
+        self.fc1 = nn.Linear(64 * 38, 128)
+        # Dropout layer, dropout rate = 0.2
+        self.dropout = nn.Dropout(0.2)
+        # fully connected layer, 5 nodes (number of classes), output shape (batch_size, 5)
+        self.fc2 = nn.Linear(128, 5)
+
+    def forward(self, x):
+        # x.shape = (batch_size, 300)
+        # reshape the tensor with shape (batch_size, 300) to (batch_size, 1, 300)
+        x = x.reshape(-1, 1, 300)
+        x = F.relu(self.conv1(x))
+        x = self.pool1(x)
+        x = F.relu(self.conv2(x))
+        x = self.pool2(x)
+        x = F.relu(self.conv3(x))
+        x = self.pool3(x)
+        x = F.relu(self.conv4(x))
+        x = self.flatten(x)
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+
+
+# define the training function and validation function
+def train_steps(loop, model, criterion, optimizer):
+    train_loss = []
+    train_acc = []
+    model.train()
+    for step_index, (X, y) in loop:
+        X, y = X.to(device), y.to(device)
+        pred = model(X)
+        loss = criterion(pred, y)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        loss = loss.item()
+        train_loss.append(loss)
+        pred_result = torch.argmax(pred, dim=1).detach().cpu().numpy()
+        y = y.detach().cpu().numpy()
+        acc = accuracy_score(y, pred_result)
+        train_acc.append(acc)
+        loop.set_postfix(loss=loss, acc=acc)
+    return {"loss": np.mean(train_loss),
+            "acc": np.mean(train_acc)}
+
+
+def test_steps(loop, model, criterion):
+    test_loss = []
+    test_acc = []
+    model.eval()
+    with torch.no_grad():
+        for step_index, (X, y) in loop:
+            X, y = X.to(device), y.to(device)
+            pred = model(X)
+            loss = criterion(pred, y).item()
+
+            test_loss.append(loss)
+            pred_result = torch.argmax(pred, dim=1).detach().cpu().numpy()
+            y = y.detach().cpu().numpy()
+            acc = accuracy_score(y, pred_result)
+            test_acc.append(acc)
+            loop.set_postfix(loss=loss, acc=acc)
+    return {"loss": np.mean(test_loss),
+            "acc": np.mean(test_acc)}
+
+
+def train_epochs(train_dataloader, test_dataloader, model, criterion, optimizer, config, writer):
+    num_epochs = config['num_epochs']
+    train_loss_ls = []
+    train_loss_acc = []
+    test_loss_ls = []
+    test_loss_acc = []
+    for epoch in range(num_epochs):
+        train_loop = tqdm(enumerate(train_dataloader), total=len(train_dataloader))
+        test_loop = tqdm(enumerate(test_dataloader), total=len(test_dataloader))
+        train_loop.set_description(f'Epoch [{epoch + 1}/{num_epochs}]')
+        test_loop.set_description(f'Epoch [{epoch + 1}/{num_epochs}]')
+
+        train_metrix = train_steps(train_loop, model, criterion, optimizer)
+        test_metrix = test_steps(test_loop, model, criterion)
+
+        train_loss_ls.append(train_metrix['loss'])
+        train_loss_acc.append(train_metrix['acc'])
+        test_loss_ls.append(test_metrix['loss'])
+        test_loss_acc.append(test_metrix['acc'])
+
+        print(f'Epoch {epoch + 1}: '
+              f'train loss: {train_metrix["loss"]}; '
+              f'train acc: {train_metrix["acc"]}; ')
+        print(f'Epoch {epoch + 1}: '
+              f'test loss: {test_metrix["loss"]}; '
+              f'test acc: {test_metrix["acc"]}')
+
+        writer.add_scalar('train/loss', train_metrix['loss'], epoch)
+        writer.add_scalar('train/accuracy', train_metrix['acc'], epoch)
+        writer.add_scalar('validation/loss', test_metrix['loss'], epoch)
+        writer.add_scalar('validation1/accuracy', test_metrix['acc'], epoch)
+
+    return {'train_loss': train_loss_ls,
+            'train_acc': train_loss_acc,
+            'test_loss': test_loss_ls,
+            'test_acc': test_loss_acc}
+
+
+def main():
+    config = {
+        'seed': 42,  # the random seed
+        'test_ratio': 0.3,  # the ratio of the test set
+        'num_epochs': 30,
+        'batch_size': 128,
+        'lr': 0.001,
+    }
+
+    # X_train,y_train is the training set
+    # X_test,y_test is the test set
+    X_train, X_test, y_train, y_test = load_data(config['test_ratio'], config['seed'])
+    # 数据集分割
+    train_dataset, test_dataset = ECGDataset(X_train, y_train), ECGDataset(X_test, y_test)
+    train_dataloader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True)
+    test_dataloader = DataLoader(test_dataset, batch_size=config['batch_size'], shuffle=False)
+
+    # define the model
+
+
+    # 主运行函数，调用model，without model, create the CNN model,use the train data to create model file
+    model = Model()
+    if os.path.exists(model_path):
+        # import the pre-trained model if it exists
+        print('Import the pre-trained model, skip the training process')
+        model.load_state_dict(torch.load(model_path))
+        model.eval()
+    else:
+        # build the CNN model
+        model = model.to(device)
+        criterion = nn.CrossEntropyLoss()
+        optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
+
+        # print the model structure
+        summary(model, (config['batch_size'], X_train.shape[1]), col_names=["input_size", "kernel_size", "output_size"],
+                verbose=2)
+
+        # define the Tensorboard SummaryWriter
+        writer = SummaryWriter(log_dir=log_dir)
+        # train and evaluate model
+        history = train_epochs(train_dataloader, test_dataloader, model, criterion, optimizer, config, writer)
+        writer.close()
+        # save the model
+        torch.save(model.state_dict(), model_path)
+        # plot the training history
+        plot_history_torch(history)
+
+    # predict the class of test data
+    y_pred = []
+    model.eval()
+    with torch.no_grad():
+        for step_index, (X, y) in enumerate(test_dataloader):
+            X, y = X.to(device), y.to(device)
+            pred = model(X)
+            pred_result = torch.argmax(pred, dim=1).detach().cpu().numpy()
+            y_pred.extend(pred_result)
+    # plot confusion matrix heat map
+    plot_heat_map(y_test, y_pred)
+
+
+if __name__ == '__main__':
+    main()

Number-Of-CNN-Layers/utils_threeLayerCNN.py

@@ -0,0 +1,172 @@
+# this file give
+import os
+
+import wfdb
+import pywt
+import seaborn
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.metrics import confusion_matrix
+from sklearn.model_selection import train_test_split
+
+# from _mian_tensorflow import Project_PATH
+Project_PATH = '../Number-Of-CNN-Layers/Two/'
+
+PICTUREPATH = Project_PATH + 'picture/'
+
+# from ECG_read_weiyuqing_without_wfdb import ECGDATAPATH
+MITBITECGDATAPATH = '../MIT-BIT/'
+
+
+# wavelet denoise preprocess using mallat algorithm
+
+def denoise(date):
+    # Wavelet db5 level9
+    coeffs = pywt.wavedec(data=date, wavelet='db5', level=9)
+    cA9, cD9, cD8, cD7, cD6, cD5, cD4, cD3, cD2, cD1 = coeffs
+    # noice del
+    threshold = (np.median(np.abs(cD1)) / 0.6745) * (np.sqrt(2 * np.log(len(cD1))))
+    # take cD 1,2 turn 0 to denoise
+    cD1.fill(0)
+    cD2.fill(0)
+    for i in range(1, len(coeffs) - 2):
+        coeffs[i] = pywt.threshold(coeffs[i], threshold)
+    # get signal that before wavelet
+    rdata = pywt.waverec(coeffs=coeffs, wavelet='db5')
+    return rdata
+
+
+# get ECG data,date type and denoise
+# this function without return
+# use append take return data to pass data
+def getDataSet(number, X_data, Y_data):
+    ecgClassSet = ['N', 'A', 'V', 'L', 'R']
+    print('get' + number + ' ECG data')
+    record = wfdb.rdrecord(MITBITECGDATAPATH + number, channel_names=['MLII'])
+    data = record.p_signal.flatten()
+    # use def denoise
+    rdata = denoise(date=data)
+
+    # get type of ECG data
+    annotation = wfdb.rdann(MITBITECGDATAPATH + number, 'atr')
+    Rlocation = annotation.sample  # annotation.sample:get R
+    Rclass = annotation.symbol
+
+    # delete stable date near R
+    start = 10
+    end = 5
+    i = start
+    j = len(annotation.symbol) - end
+
+    # get NAVLR type data
+    while i < j:
+        try:
+            lable = ecgClassSet.index(Rclass[i])
+            x_train = rdata[Rlocation[i] - 99:Rlocation[i] + 201]
+            X_data.append(x_train)
+            Y_data.append(lable)
+            i += 1
+        except ValueError:
+            i += 1
+    return
+
+
+# RATIO=The ratio of the training set to the test set，usually is 0.3
+def loadData(ratio, random_seed):
+    numberSet = ['100', '101', '103', '105', '106', '107', '108', '109', '111', '112', '113', '114', '115',
+                 '116', '117', '119', '121', '122', '123', '124', '200', '201', '202', '203', '205', '208',
+                 '210', '212', '213', '214', '215', '217', '219', '220', '221', '222', '223', '228', '230',
+                 '231', '232', '233', '234']
+    dataSet = []
+    labelSet = []
+    for n in numberSet:
+        getDataSet(n, dataSet, labelSet)
+
+    # dataSet = np.array(dataSet).reshape(-1, 300)
+    # labelSet = np.array(labelSet).reshape(-1, 1)
+
+    # reshape the data and split the dataset
+    dataSet = np.array(dataSet).reshape(-1, 300)
+    lableSet = np.array(labelSet).reshape(-1)
+    X_train, X_test, y_train, y_test = train_test_split(dataSet, lableSet, test_size=ratio, random_state=random_seed)
+    # print(X_train, X_test, y_train, y_test)
+    return X_train, X_test, y_train, y_test
+
+    # train_ds = np.hstack((dataSet, labelSet))
+    # np.random.shuffle(train_ds)
+    #
+    # X = train_ds[:, :300].reshape(-1, 300, 1)
+    # Y = train_ds[:, 300]
+    # # cut test and train data from ECG data
+    # shuffle_index = np.random.permutation(len(X))
+    # test_length = int(RATIO * len(shuffle_index))
+    # test_index = shuffle_index[:test_length]
+    # train_index = shuffle_index[test_index:]
+    # X_test, Y_test = X[test_index:], Y[test_index]
+    # X_train, Y_train = X[train_index], Y[train_index]
+    # return X_train, Y_test, X_test, Y_train
+
+
+# confusion matrix
+def plot_heat_map(y_test, y_pred):
+    con_mat = confusion_matrix(y_test, y_pred)
+    # normalize
+    # con_mat_norm = con_mat.astype('float') / con_mat.sum(axis=1)[:, np.newaxis]
+    # con_mat_norm = np.around(con_mat_norm, decimals=2)
+    # os.remove(PICTUREPATH+'confusion_matrix.png')
+    # plot
+    plt.figure(figsize=(8, 8))
+    seaborn.heatmap(con_mat, annot=True, fmt='.20g', cmap='Blues')
+    plt.ylim(0, 5)
+    plt.xlabel('Predicted labels')
+    plt.ylabel('True labels')
+    plt.title('Confusion Matrix')
+    plt.savefig(PICTUREPATH + 'confusion_matrix.png')
+    plt.show()
+
+
+def plot_history_tf(history):
+    # os.remove(PICTUREPATH+'accuracy.png')
+    # os.remove(PICTUREPATH+'loss.png')
+
+    plt.figure(figsize=(8, 8))
+    plt.plot(history.history['accuracy'])
+    plt.plot(history.history['val_accuracy'])
+    plt.title('Model Accuracy')
+    plt.ylabel('Accuracy')
+    plt.xlabel('Epoch')
+    plt.legend(['Train', 'Test'], loc='upper left')
+    plt.savefig(PICTUREPATH + 'accuracy.png')
+    plt.show()
+
+    plt.figure(figsize=(8, 8))
+    plt.plot(history.history['loss'])
+    plt.plot(history.history['val_loss'])
+    plt.title('Model Loss')
+    plt.ylabel('Loss')
+    plt.xlabel('Epoch')
+    plt.legend(['Train', 'Test'], loc='upper left')
+    plt.savefig(PICTUREPATH + 'loss.png')
+    plt.show()
+
+
+def plot_history_torch(history):
+    plt.figure(figsize=(8, 8))
+    plt.plot(history['train_acc'])
+    plt.plot(history['test_acc'])
+    plt.title('Model Accuracy')
+    plt.ylabel('Accuracy')
+    plt.xlabel('Epoch')
+    plt.legend(['Train', 'Test'], loc='upper left')
+    plt.savefig(PICTUREPATH + 'accuracy.png')
+    plt.show()
+
+    plt.figure(figsize=(8, 8))
+    plt.plot(history['train_loss'])
+    plt.plot(history['test_loss'])
+    plt.title('Model Loss')
+    plt.ylabel('Loss')
+    plt.xlabel('Epoch')
+    plt.legend(['Train', 'Test'], loc='upper left')
+    plt.savefig(PICTUREPATH + 'loss.png')
+    plt.show()