import torch import torch.nn as nn from torchvision import transforms, datasets, utils import matplotlib.pyplot as plt import numpy as np import torch.optim as optim #from model_improved_alexnet import ImprovedAlexNet#调用模型 #from model_improved_alexnet_plus import ImprovedAlexNet from Se_AlexNet import ImprovedAlexNet #from model_alexnet import ImprovedAlexNet import os import json import time # 配置环境 #device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") device = torch.device("cpu") print(device) # 定义数据转换器 data_transform = { "train": transforms.Compose([transforms.RandomResizedCrop(224),#调整图像尺寸为224*224 transforms.RandomHorizontalFlip(),#数据增强，对图像进行随机水平翻转 transforms.ToTensor(),#将图像数据转化为张量 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),#数据标准化 "val": transforms.Compose([transforms.Resize((224, 224)), # cannot 224, must (224, 224) transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])} # 数据集路径 data_root = os.path.abspath(os.path.join(os.getcwd())) image_path = data_root + "/data_set/sn_data/" train_dataset = datasets.ImageFolder(root=image_path + "/train", transform=data_transform["train"]) #print(train_dataset) train_num = len(train_dataset)#获取训练数据集的样本数量 # 类别列表：将训练数据集的类别索引信息保存到一个 JSON 文件中，方便在以后的程序中快速查找类别与索引之间的对应关系。 class_list = train_dataset.class_to_idx #print(class_list) cla_dict = dict((val, key) for key, val in class_list.items()) #print(cla_dict) # write dict into json file json_str = json.dumps(cla_dict, indent=4) #print(json_str) with open('class_indices.json', 'w') as json_file: json_file.write(json_str) # 设置数据加载器 batch_size = 32#数据集的批量大小为 32 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True,#在每个epoch开始时打乱数据顺序，可以增加数据的随机性，有助于模型的泛化能力 num_workers=0) validate_dataset = datasets.ImageFolder(root=image_path + "/val", transform=data_transform["val"])#加载验证数据集 val_num = len(validate_dataset)#获取验证数据集的样本数量 validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=4, shuffle=True, num_workers=0)#同上 # 网络配置 net = ImprovedAlexNet(num_classes=11)#定义了net网络模型，模型实例化 #预训练权重 #model_weight_path = "models/alexnet.pth" #net.load_state_dict(torch.load(model_weight_path, map_location=torch.device('cpu')), strict=False) #model = torch.load('path_to_model.pth', map_location=torch.device('cpu')) #model = torch.load('model.pth', map_location=torch.device('cpu')) net.to(device) # 损失函数 loss_function = nn.CrossEntropyLoss() # pata = list(net.parameters()) optimizer = optim.Adam(net.parameters(), lr=0.0002) save_path = './models/model.pth' best_acc = 0.0#初始化最佳验证准确率，用于保存模型训练过程中的最佳准确率 from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix import seaborn as sns import torchmetrics import matplotlib.pyplot as plt plt.rcParams["font.sans-serif"]=["SimHei"] #设置字体 plt.rcParams["axes.unicode_minus"]=False #该语句解决图像中的“-”负号的乱码问题,设置负号显示为正常的减号符号，而不是方块 # 初始化metrics 初始化指标， mdmc_average='samplewise', #参数 average 表示 Precision 的计算方式，'macro' 表示计算每个类别的 Precision，并对所有类别的 Precision 取平均值 train_accuracy = torchmetrics.Accuracy(num_classes=len(class_list), average='macro', task='multiclass').to(device) val_accuracy = torchmetrics.Accuracy(num_classes=len(class_list), average='macro', task='multiclass').to(device) train_precision = torchmetrics.Precision(num_classes=len(class_list), average='macro', task='multiclass').to(device)#取平均值 val_precision = torchmetrics.Precision(num_classes=len(class_list), average='macro', task='multiclass').to(device) train_recall = torchmetrics.Recall(num_classes=len(class_list), average='macro', task='multiclass').to(device) val_recall = torchmetrics.Recall(num_classes=len(class_list), average='macro', task='multiclass').to(device) train_f1 = torchmetrics.F1Score(num_classes=len(class_list), average='macro', task='multiclass').to(device) val_f1 = torchmetrics.F1Score(num_classes=len(class_list), average='macro', task='multiclass').to(device) # 用于绘图的列表，记录每次迭代的数据 epoch_count = [] train_losses = [] val_losses = [] train_accuracies = [] val_accuracies = [] train_precisions = [] val_precisions = [] train_recalls = [] val_recalls = [] train_f1s = [] val_f1s = [] # 初始化记录最优情况下的性能指标 best_precision = 0.0 best_recall = 0.0 best_f1 = 0.0 best_confusion_matrix = None best_epoch = 0 epochs=70 # 训练模型 for epoch in range(epochs): # train net.train() running_loss = 0.0 t1 = time.perf_counter()#记录当前时间，用于计算每个epoch的训练时间 for step, data in enumerate(train_loader, start=0): images, labels = data#获取数据和标签 optimizer.zero_grad()#清除之前保存的梯度，以便进行新一轮的梯度计算和优化 outputs = net(images.to(device))#输出 loss = loss_function(outputs, labels.to(device))#计算输出outputs与真实标签labels之间的损失值 loss.backward() optimizer.step() # print statistics running_loss += loss.item()#累加当前batch的损失值 train_accuracy(outputs.to(device), labels.to(device))#计算准确度 train_precision(outputs.to(device), labels.to(device))#计算精确度 train_recall(outputs.to(device), labels.to(device))#计算召回率 train_f1(outputs.to(device), labels.to(device))#计算F1分数 # print train process rate = (step + 1) / len(train_loader)#显示进度条 a = "*" * int(rate * 50) b = "." * int((1 - rate) * 50) print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="") print() #print(time.perf_counter()-t1) # validate net.eval() val_accuracy.reset() val_precision.reset() val_recall.reset() val_f1.reset() val_loss = 0.0 acc = 0.0 # accumulate accurate number / epoch all_predicts = [] all_labels = [] with torch.no_grad(): for val_data in validate_loader: val_images, val_labels = val_data outputs = net(val_images.to(device)) predict_y = torch.max(outputs, dim=1)[1] acc += (predict_y == val_labels.to(device)).sum().item() val_loss = loss_function(outputs.to(device), val_labels.to(device)) val_accuracy(outputs.to(device), val_labels.to(device)) val_precision(outputs.to(device), val_labels.to(device)) val_recall(outputs.to(device), val_labels.to(device)) val_f1(outputs.to(device), val_labels.to(device)) # 保存所有的预测和标签，用于后续的性能评估 all_predicts.extend(predict_y.cpu().numpy()) all_labels.extend(val_labels.cpu().numpy()) # 计算性能指标 val_accurate = acc / val_num precision = precision_score(all_labels, all_predicts, average='macro') recall = recall_score(all_labels, all_predicts, average='macro') f1 = f1_score(all_labels, all_predicts, average='macro') confusion = confusion_matrix(all_labels, all_predicts)#混淆矩阵 # 检查是否是最佳性能，如果是，更新最佳性能指标和保存模型 if val_accurate > best_acc: best_acc = val_accurate best_precision = precision best_recall = recall best_f1 = f1 best_confusion_matrix = confusion best_epoch = epoch torch.save(net.state_dict(), save_path) # 计算平均损失和准确率 epoch_loss = running_loss / len(train_loader) epoch_val_loss = val_loss.item() / len(validate_loader) epoch_train_accuracy = train_accuracy.compute() epoch_val_accuracy = val_accuracy.compute() epoch_train_precision = train_precision.compute() epoch_val_precision = val_precision.compute() epoch_train_recall = train_recall.compute() epoch_val_recall = val_recall.compute() epoch_train_f1 = train_f1.compute() epoch_val_f1 = val_f1.compute() # 记录历史数据以便绘图 epoch_count.append(epoch + 1) train_losses.append(epoch_loss) val_losses.append(epoch_val_loss) train_accuracies.append(epoch_train_accuracy.item()) val_accuracies.append(epoch_val_accuracy.item()) train_precisions.append(epoch_train_precision.item()) val_precisions.append(epoch_val_precision.item()) train_recalls.append(epoch_train_recall.item()) val_recalls.append(epoch_val_recall.item()) train_f1s.append(epoch_train_f1.item()) val_f1s.append(epoch_val_f1.item()) # 打印每个epoch的指标 print(f'Epoch {epoch + 1}/{epochs}, Train Loss: {epoch_loss:.4f}, Validation Loss: {epoch_val_loss:.4f}, ' f'Train Accuracy: {epoch_train_accuracy:.4f}, Validation Accuracy: {epoch_val_accuracy:.4f}, ' f'Train Precision: {epoch_train_precision:.4f}, Validation Precision: {epoch_val_precision:.4f}, ' f'Train Recall: {epoch_train_recall:.4f}, Validation Recall: {epoch_val_recall:.4f}, ' f'Train F1: {epoch_train_f1:.4f}, Validation F1: {epoch_val_f1:.4f}') print('Finished Training') # 保存最佳性能指标到txt with open('./logs/model_performance.txt', 'w', encoding='utf-8') as fb: fb.write(f'Accuracy: {best_acc}\n') fb.write(f'Precision: {best_precision}\n') fb.write(f'Recall: {best_recall}\n') fb.write(f'F1-Score: {best_f1}\n') fb.write(f'Confusion Matrix: \n{best_confusion_matrix}\n') # 绘制训练和验证损失图 plt.figure(figsize=(12, 8)) plt.subplot(2, 3, 1) plt.plot(epoch_count, train_losses, label='Train Loss') plt.plot(epoch_count, val_losses, label='Test Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.title('Training and Test Loss') plt.legend() # 绘制训练和验证准确率图 plt.subplot(2, 3, 2) plt.plot(epoch_count, train_accuracies, label='Train Accuracy') plt.plot(epoch_count, val_accuracies, label='Test Accuracy') plt.xlabel('Epoch') plt.ylabel('Accuracy') plt.title('Training and Test Accuracy') plt.legend() # 绘制训练和验证精确度图 plt.subplot(2, 3, 3) plt.plot(epoch_count, train_precisions, label='Train Precision') plt.plot(epoch_count, val_precisions, label='Test Precision') plt.xlabel('Epoch') plt.ylabel('Precision') plt.title('Training and Test Precision') plt.legend()#为图表添加说明性文本或标签的方法，以便在数据可视化中更好地传达数据信息 # 绘制训练和验证召回率图 plt.subplot(2, 3, 4) plt.plot(epoch_count, train_recalls, label='Train Recall') plt.plot(epoch_count, val_recalls, label='Test Recall') plt.xlabel('Epoch') plt.ylabel('Recall') plt.title('Training and Test Recall') plt.legend() # 绘制训练和验证F1分数图 plt.subplot(2, 3, 5) plt.plot(epoch_count, train_f1s, label='Train F1 Score') plt.plot(epoch_count, val_f1s, label='Test F1 Score') plt.xlabel('Epoch') plt.ylabel('F1 Score') plt.title('Training and Test F1 Score') plt.legend() plt.tight_layout() plt.show() if best_confusion_matrix is not None: # 设置图的大小 plt.figure(figsize=(10, 8)) # 绘制混淆矩阵的热图，Seaborn库中的heatmap函数 sns.heatmap(best_confusion_matrix, annot=True, fmt="d", cmap='Blues', xticklabels=cla_dict.values(), yticklabels=cla_dict.values()) # 使用类别名称作为标签 # 设置图的标题和坐标轴标签 plt.title('Confusion Matrix at Best Performance') plt.xlabel('Predicted Label') plt.ylabel('True Label') # 显示图表 plt.show() else: print("Best confusion matrix is not available.") # 写出准确率 with open('./logs/model.txt', 'w', encoding='utf-8') as fb: fb.write(str(best_acc))