diff --git a/Number-Of-CNN-Layers/_main_pytorch_threeLayerCNN.py b/Number-Of-CNN-Layers/_main_pytorch_threeLayerCNN.py new file mode 100644 index 0000000..5afe3ad --- /dev/null +++ b/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() diff --git a/Number-Of-CNN-Layers/utils_threeLayerCNN.py b/Number-Of-CNN-Layers/utils_threeLayerCNN.py new file mode 100644 index 0000000..d06c181 --- /dev/null +++ b/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()