CNN 基于 DAS 数据的六类故障分类模型
coding=UTF-8\nimport sys\nimport os\n\n\nclass Logger(object):\n def init(self, filename='default.log', stream=sys.stdout):\n self.terminal = stream\n self.log = open(filename, 'w')\n\n def write(self, message):\n self.terminal.write(message)\n self.log.write(message)\n\n def flush(self):\n pass\n\n\nsys.stdout = Logger('result.log', sys.stdout)\nimport argparse\nimport numpy as np\nimport time\nfrom sklearn.metrics import accuracy_score, confusion_matrix\nfrom models import CNN\nfrom mydataset import MyDataset\nimport torch\nfrom torch.utils.data import DataLoader\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport seaborn as sns # µ¼Èë°ü\nimport matplotlib.font_manager as fm\nfrom tqdm import *\nimport csv\n\n\ndef test(model, dataset, criterion):\n model.eval()\n total_batch_num = 0\n val_loss = 0\n prediction = []\n labels = []\n feature_list = torch.tensor([]) # test×ÜÊý£¬ÌØÕ÷Êý\n if torch.cuda.is_available():\n feature_list = feature_list.cuda()\n for (step, i) in enumerate(dataset):\n total_batch_num = total_batch_num + 1\n batch_x = i['data']\n batch_y = i['label']\n batch_x = torch.unsqueeze(batch_x, dim=1) # (50, 1, 10000, 12)\n batch_x = batch_x.float()\n if torch.cuda.is_available():\n batch_x = batch_x.cuda()\n batch_y = batch_y.cuda()\n feature, probs = model(batch_x) # feature 20,400\n batch_label = batch_y.unsqueeze(1).float()\n feature_label = torch.cat((feature, batch_label), dim=1)\n feature_list = torch.cat((feature_list, feature_label), dim=0) # feature_list\n loss = criterion(probs, batch_y)\n , pred = torch.max(probs, dim=1)\n predi = pred.tolist()\n label = batch_y.tolist()\n val_loss += loss.item()\n prediction.extend(predi)\n labels.extend(label)\n accuracy = accuracy_score(labels, prediction)\n C = confusion_matrix(labels, prediction)\n return accuracy, val_loss / total_batch_num, feature_list, C\n\n\ndef train(model, train_x, train_y, optimizer, criterion):\n model.train()\n model.zero_grad()\n ,probs = model(train_x)\n loss = criterion(probs, train_y)\n , pred = torch.max(probs, dim=1)\n labels = train_y.tolist()\n predi = pred.tolist()\n loss.backward()\n optimizer.step()\n return labels, predi, loss.item()\n\n\ndef draw(train_acc, train_loss, test_acc, test_loss):\n x1 = range(len(train_acc))\n x2 = range(len(train_loss))\n y1 = train_acc\n y2 = train_loss\n y3 = test_acc\n y4 = test_loss\n plt.subplot(2, 1, 1)\n plt.plot(x1, y1, 'o-', label="train", color='b')\n plt.plot(x1, y3, 'o-', label="test", color='r')\n plt.legend(loc='upper left')\n plt.title('accuracy & NAR vs. epochs')\n plt.ylabel('accuracy')\n plt.subplot(2, 1, 2)\n plt.plot(x2, y2, '.-', label="train", color='b')\n plt.plot(x2, y4, '.-', label="test", color='r')\n plt.legend(loc='upper left')\n plt.xlabel('loss vs. epochs')\n plt.ylabel('loss')\n plt.savefig("accuracy_loss.jpg")\n plt.show()\n\n\ndef draw_result(C):\n fig = plt.figure()\n ax = fig.add_subplot(111)\n df = pd.DataFrame(C)\n # f1 = fm.fontProperties('Times New Roman', size=15)\n sns.heatmap(df, fmt='g', annot=True,\n annot_kws={'size': 10},\n xticklabels=['1', '2', '3', '4', '5', '6'],\n yticklabels=['1', '2', '3', '4', '5', '6'],\n cmap='Blues')\n ax.set_xlabel('Predicted label') # fontProperties=f1) # xÖá\n ax.set_ylabel('True label') # fontProperties=f1) # yÖá\n plt.savefig('./CNN_confusion_matrix.jpg')\n plt.show()\n Acc = (C[0][0] + C[1][1] + C[2][2] + C[3][3] + C[4][4] + C[5][5]) / sum(C[0] + C[1] + C[2] + C[3] + C[4] + C[5])\n print('acc: %.3f' % Acc)\n lie_he = sum(C, 1) - 1\n for i in range(1, 7):\n Precision = C[i - 1][i - 1] / lie_he[i - 1]\n NAR = (sum(C[i - 1]) - C[i - 1][i - 1]) / sum(C[i - 1])\n F1_score = 2 * C[i - 1][i - 1] / (lie_he[i - 1] + sum(C[i - 1]))\n print('precision%d: %.3f' % (i, Precision))\n print('NAR%d: %.3f' % (i, NAR))\n print('F1_score%d: %.3f' % (i, F1_score))\n\n\ndef main(args):\n # »®·ÖѵÁ·¼¯¡¢²âÊÔ¼¯\n train_dataset = MyDataset(args.root, args.txtpath, transform=None)\n train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)\n test_dataset = MyDataset(args.root2, args.txtpath2, transform=None)\n test_loader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)\n\n models = {"CNN": CNN}\n\n\n\n # ¶þÕßÈ¡ÆäÒ»\n #model = torch.load('./modelpth/finalmodel.pth')#,map_location=torch.device('cpu')) # ¼ÓÔØ²ÎÊý\n model = modelsargs.model\n\n if torch.cuda.is_available():\n model = model.cuda()\n\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)\n # optimizer.load_state_dict(torch.load('save_optim.pt'))\n criterion = nn.CrossEntropyLoss()\n batches_per_epoch = int(len(train_dataset) / args.batch_size)\n train_loss_list = []\n train_acc_list = []\n test_loss_list = []\n test_acc_list = []\n train_time = 0\n model_save_path = './modelpth/'\n os.makedirs(model_save_path, exist_ok=True)\n for epoch in range(args.epochs):\n tic = time.time()\n train_predict = []\n train_label = []\n runloss = 0\n for (cnt, i) in enumerate(tqdm(train_loader)):\n batch_x = i['data']\n batch_y = i['label']\n batch_x = torch.unsqueeze(batch_x, dim=1)\n batch_x = batch_x.float()\n if torch.cuda.is_available():\n batch_x = batch_x.cuda()\n batch_y = batch_y.cuda()\n tlabels, tpredi, tloss = train(model, batch_x, batch_y, optimizer, criterion)\n runloss = runloss + tloss\n train_label.extend(tlabels)\n train_predict.extend(tpredi)\n per_epoch_train_time = time.time() - tic\n train_time = per_epoch_train_time + train_time\n taccuracy = accuracy_score(train_label, train_predict)\n train_acc_list.append(taccuracy)\n loss = runloss / batches_per_epoch\n train_loss_list.append(loss)\n acc_score, loss_score, feature_list, C = test(model, test_loader, criterion)\n print("Epoch %d Val_accuracy %.3f Val_loss %.3f" % (epoch, acc_score, loss_score))\n with open('log.csv', 'a+')as f:\n f.write(f'{epoch}, {round(acc_score, 4)}, {round(loss_score, 4)}\n')\n f.close()\n test_acc_list.append(acc_score)\n test_loss_list.append(loss_score)\n if epoch % 10 == 0:\n torch.save(model, f'./modelpth/{epoch}.pth')\n if epoch == args.epochs - 1: # ѵÁ·µ½×îºóÒ»ÂÖ\n torch.save(model, './modelpth/finalmodel.pth')\n the_model = torch.load('./modelpth/finalmodel.pth')\n acc_score, loss_score, NAR, featucre_list = test(the_model, test_loader, criterion)\n #print("Epoch %d Train_accuracy %.3f Train_loss %.3f Val_accuracy %.3f Val_loss %.3f NAR %.3f " % (epoch, taccuracy, loss, acc_score, loss_score, NAR))\n print(f'Epoch {epoch} Train_accuracy {taccuracy} Train_loss {loss} Val_accuracy {acc_score} Val_loss {loss_score} NAR {NAR}')\n feature_list = feature_list.detach().cpu().numpy()\n np.savetxt('feature_data.csv', feature_list, delimiter=',') # ±£´æÌØÕ÷ÏòÁ¿\n draw_result(C)\n print('train totally using %.3f seconds ', train_time)\n draw(train_acc_list, train_loss_list, test_acc_list, test_loss_list)\n\n\nif name == "main":\n parser = argparse.ArgumentParser(description="CNN fou classification")\n '''save model'''\n parser.add_argument("--save", type=str, default="__", help="path to save model")\n '''model parameters'''\n rootpath = 'das_data'\n parser.add_argument("--root", type=str, default=rootpath + '/train', help="rootpath of traindata")\n parser.add_argument("--root2", type=str, default=rootpath + '/test', help="rootpath of valdata")\n parser.add_argument("--txtpath", type=str, default=rootpath + '/train/label.txt', help="path of train_list")\n parser.add_argument("--txtpath2", type=str, default=rootpath + '/test/label.txt', help="path pf val_list")\n parser.add_argument("--model", type=str, default="CNN", help="type of model to use for classification")\n parser.add_argument("--lr", type=float, default=0.01, help="learning rate")\n parser.add_argument("--epochs", type=int, default=50 , help="number of training epochs")\n parser.add_argument("--batch_size", type=int, default=100, help="batch size")\n my_args = parser.parse_args()\n\n main(my_args)
原文地址: https://www.cveoy.top/t/topic/pUj0 著作权归作者所有。请勿转载和采集!