web3-202226701039/ChexnetTrainer(2)(4).py

import os
import numpy as np
import time
import sys
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import torch.nn.functional as tfunc
import tqdm
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau
import torch.nn.functional as func
from tqdm import tqdm
from sklearn.metrics import roc_auc_score,roc_curve,auc,f1_score
from DensenetModels import DenseNet121, DenseNet169, \
    DenseNet201,ResNet50,FocalLoss  # 引入不同版本的DenseNet
from DatasetGenerator import DatasetGenerator  # 引入自定义的数据集生成器


# 定义一个用于训练、验证、测试DenseNet模型的类
class ChexnetTrainer():
    # 训练网络的主函数
    def train(pathDirData, pathFileTrain, pathFileVal, nnArchitecture,
              nnIsTrained, nnClassCount, trBatchSize, trMaxEpoch, transResize,
              transCrop, launchTimestamp, checkpoint):

        # 根据选择的模型架构来初始化不同的DenseNet模型
        if nnArchitecture == 'DENSE-NET-121':
            model = DenseNet121(nnClassCount,
                                nnIsTrained).cuda()  # 初始化DenseNet121
        elif nnArchitecture == 'DENSE-NET-169':
            model = DenseNet169(nnClassCount,
                                nnIsTrained).cuda()  # 初始化DenseNet169
        elif nnArchitecture == 'DENSE-NET-201':
            model = DenseNet201(nnClassCount,
                             nnIsTrained).cuda()  # 初始化DenseNet201
        elif nnArchitecture == 'RESNET-50':
            model = ResNet50(nnClassCount, nnIsTrained).cuda()
        else:
            raise ValueError(
                f"Unknown architecture: {nnArchitecture}. Please choose from 'DENSE-NET-121', 'DENSE-NET-169', 'DENSE-NET-201', 'RESNET-50'.")
        model = model.cuda()  # 将模型加载到GPU上

        # 数据预处理，包含随机裁剪、水平翻转、归一化等
        normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                         [0.229, 0.224, 0.225])
        transformList = [
            transforms.RandomResizedCrop(transCrop),  # 随机裁剪到指定大小
            transforms.RandomHorizontalFlip(),  # 随机水平翻转
            transforms.ToTensor(),  # 转换为张量
            normalize  # 归一化
        ]
        transformSequence = transforms.Compose(transformList)  # 将这些变换组成序列

        # 创建训练和验证集的数据加载器
        datasetTrain = DatasetGenerator(pathImageDirectory=pathDirData,
                                        pathDatasetFile=pathFileTrain,
                                        transform=transformSequence)
        #pos_weight = datasetTrain.calculate_pos_weights()  # 计算 pos_weight
        datasetVal = DatasetGenerator(pathImageDirectory=pathDirData,
                                      pathDatasetFile=pathFileVal,
                                      transform=transformSequence)
        dataLoaderTrain = DataLoader(dataset=datasetTrain,
                                     batch_size=trBatchSize, shuffle=True,
                                     num_workers=8, pin_memory=True)  # 训练集
        dataLoaderVal = DataLoader(dataset=datasetVal, batch_size=trBatchSize,
                                   shuffle=False, num_workers=8,
                                   pin_memory=True)  # 验证集

        # 设置优化器和学习率调度器
        optimizer = optim.Adam(model.parameters(), lr=0.0001,
                               betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)
        scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=5,
                                      mode='min')  # 当损失不再下降时，减少学习率

        # 使用多标签分类的损失函数
        # 设置每个类别的权重，基于提供的AUC-ROC曲线
        # class_weights = torch.tensor(
        #     [1.39, 0.73, 1.33, 1.62, 1.32, 1.41, 1.54, 1.27, 1.59, 1.25, 1.28,
        #      1.27, 1.48, 1.18], dtype=torch.float).cuda()
        # 使用加权的 BCEWithLogitsLoss 作为损失函数
        #loss = torch.nn.BCEWithLogitsLoss(pos_weight=class_weights)
        #loss = torch.nn.MultiLabelSoftMarginLoss()
        # 使用Focal Loss作为损失函数
        # loss = FocalLoss(alpha=1, gamma=2, logits=True)
        # launchTimestamp += str(loss)
        # launchTimestamp += 'delete'
        class_weights = torch.tensor([
    1.1762,  # Atelectasis
    0.6735,  # Cardiomegaly
    0.9410,  # Effusion
    1.6680,  # Infiltration
    0.9699,  # Mass
    1.1950,  # Nodule
    1.7584,  # Pneumonia
    0.6859,  # Pneumothorax
    1.6683,  # Consolidation
    0.7744,  # Edema
    0.4625,  # Emphysema
    0.7385,  # Fibrosis
    1.3764,  # Pleural_Thickening
    0.1758   # Hernia
], dtype=torch.float32).cuda()
        loss = FocalLoss(alpha=class_weights, gamma=2, logits=True)


        # 加载检查点文件（如果存在），继续训练
        #checkpoint = 'm-29102024-093913MultiLabelSoftMarginLoss()delete.pth.tar'  # 测试时写死的文件名
        # if checkpoint != None:
        #     modelCheckpoint = torch.load(checkpoint)
        #     model.load_state_dict(modelCheckpoint['state_dict'])
        #     optimizer.load_state_dict(modelCheckpoint['optimizer'])

        lossMIN = 100000  # 记录最小损失值
        train_f1_scores, val_f1_scores = [], []
        # 训练循环，遍历指定的epoch数
        for epochID in range(trMaxEpoch):

            # 获取当前时间戳，记录每个epoch的开始时间
            timestampTime = time.strftime("%H%M%S")
            timestampDate = time.strftime("%d%m%Y")
            timestampSTART = timestampDate + '-' + timestampTime

            # 训练一个epoch
            ChexnetTrainer.epochTrain(model, dataLoaderTrain, optimizer,
                                      scheduler, trMaxEpoch, nnClassCount, loss,train_f1_scores)
            # 验证一个epoch
            lossVal, losstensor,val_f1 = ChexnetTrainer.epochVal(model, dataLoaderVal,optimizer, scheduler,
                                                          trMaxEpoch,nnClassCount, loss)
            val_f1_scores.append(val_f1)
            # 获取每个epoch结束时的时间戳
            timestampTime = time.strftime("%H%M%S")
            timestampDate = time.strftime("%d%m%Y")
            timestampEND = timestampDate + '-' + timestampTime

            # 使用调度器调整学习率
            scheduler.step(losstensor.item())

            # 保存当前最佳模型
            if lossVal < lossMIN:
                lossMIN = lossVal
                torch.save(
                    {'epoch': epochID + 1, 'state_dict': model.state_dict(),
                     'best_loss': lossMIN, 'optimizer': optimizer.state_dict()},
                    'm-' + launchTimestamp + '.pth.tar')
                print('Epoch [' + str(
                    epochID + 1) + '] [save] [' + timestampEND + '] loss= ' + str(
                    lossVal))
            else:
                print('Epoch [' + str(
                    epochID + 1) + '] [----] [' + timestampEND + '] loss= ' + str(
                    lossVal))
        plt.plot(train_f1_scores, label="Train F1-Score")
        plt.plot(val_f1_scores, label="Val F1-Score")
        plt.xlabel("Epoch")
        plt.ylabel("F1 Score")
        plt.title("F1 Score per Epoch")
        plt.legend()
        #plt.savefig("f1score.png")
        #plt.show()
    # 训练每个epoch的具体过程
    def epochTrain(model, dataLoader, optimizer, scheduler, epochMax,
                   classCount, loss,f1_scores):
        model.train()
        all_targets = []
        all_preds = []

        for batchID, (input, target) in enumerate(tqdm(dataLoader)):
            target = target.cuda()
            input = input.cuda()

            varInput = torch.autograd.Variable(input)
            varTarget = torch.autograd.Variable(target)
            varOutput = model(varInput)

            lossvalue = loss(varOutput, varTarget)

            optimizer.zero_grad()
            lossvalue.backward()
            optimizer.step()

            pred = torch.sigmoid(varOutput).cpu().data.numpy() > 0.5
            all_targets.extend(target.cpu().numpy())
            all_preds.extend(pred)


        f1 = f1_score(np.array(all_targets), np.array(all_preds), average="macro")
        f1_scores.append(f1)


    # 验证每个epoch的具体过程
    def epochVal(model, dataLoader, optimizer, scheduler, epochMax, classCount,
                 loss):
        model.eval()  # 设置模型为评估模式

        lossVal = 0
        lossValNorm = 0
        losstensorMean = 0
        all_targets = []
        all_preds = []

        for i, (input, target) in enumerate(dataLoader):  # 遍历每个批次
            target = target.cuda()
            input = input.cuda()
            with torch.no_grad():  # 禁用梯度计算，节省内存
                varInput = torch.autograd.Variable(input)
                varTarget = torch.autograd.Variable(target)
                varOutput = model(varInput)  # 前向传播

            losstensor = loss(varOutput, varTarget)  # 计算损失
            losstensorMean += losstensor  # 累积损失
            lossVal += losstensor.item()  # 累积损失值
            lossValNorm += 1  # 记录批次数量

            pred = torch.sigmoid(varOutput).cpu().data.numpy() > 0.5
            all_targets.extend(target.cpu().numpy())
            all_preds.extend(pred)
        f1 = f1_score(np.array(all_targets), np.array(all_preds),
                      average="macro")
        outLoss = lossVal / lossValNorm  # 计算平均损失
        losstensorMean = losstensorMean / lossValNorm  # 平均损失张量

        return outLoss, losstensorMean,f1  # 返回验证损失和损失张量均值

    # 计算AUROC（AUC-ROC曲线下的面积）
    def computeAUROC(dataGT, dataPRED, classCount, plot_roc_curve=False,
                     class_names=None):
        outAUROC = []
        datanpGT = dataGT.cpu().numpy()  # 将ground truth转换为numpy格式
        datanpPRED = dataPRED.cpu().numpy()  # 将预测结果转换为numpy格式

        if plot_roc_curve and class_names is None:
            class_names = [f"Class {i + 1}" for i in range(classCount)]

        # 针对每个类别计算ROC AUC分数
        plt.figure(figsize=(12, 8))
        for i in range(classCount):
            # 计算当前类别的ROC AUC分数
            outAUROC.append(roc_auc_score(datanpGT[:, i], datanpPRED[:, i]))
            if plot_roc_curve:
                # 计算 ROC 曲线的点
                fpr, tpr, _ = roc_curve(datanpGT[:, i], datanpPRED[:, i])
                roc_auc = auc(fpr, tpr)
                plt.plot(fpr, tpr, lw=2,
                         label=f'{class_names[i]} (area = {roc_auc:.2f})')

        if plot_roc_curve:
            plt.plot([0, 1], [0, 1], color='navy', linestyle='--')  # 绘制随机猜测线
            plt.xlim([0.0, 1.0])
            plt.ylim([0.0, 1.05])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver Operating Characteristic (ROC) Curves')
            plt.legend(loc="lower right")
            plt.savefig("aucroc.png")
            #plt.show()

        return outAUROC  # 返回每个类别的AUROC值

    # 测试模型
    def test(pathDirData, pathFileTest, pathModel, nnArchitecture, nnClassCount,
             nnIsTrained, trBatchSize, transResize, transCrop, launchTimeStamp):
        CLASS_NAMES = ['Atelectasis', 'Cardiomegaly', 'Effusion',
                       'Infiltration', 'Mass', 'Nodule', 'Pneumonia',
                       'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema',
                       'Fibrosis', 'Pleural_Thickening', 'Hernia']

        cudnn.benchmark = True  # 加速卷积操作

        # 根据架构选择相应的DenseNet模型
        if nnArchitecture == 'DENSE-NET-121':
            model = DenseNet121(nnClassCount, nnIsTrained).cuda()
        elif nnArchitecture == 'DENSE-NET-169':
            model = DenseNet169(nnClassCount, nnIsTrained).cuda()
        elif nnArchitecture == 'DENSE-NET-201':
            model = DenseNet201(nnClassCount, nnIsTrained).cuda()
        elif nnArchitecture == 'RESNET-50':
            model = ResNet50(nnClassCount, nnIsTrained).cuda()

        modelCheckpoint = torch.load(pathModel)  # 加载模型
        model.load_state_dict(modelCheckpoint['state_dict'])  # 载入训练好的参数

        model = model.cuda()  # 将模型加载到GPU上
        model.eval()  # 设置为评估模式

        # 定义数据预处理
        normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                         [0.229, 0.224, 0.225])
        transformList = [
            transforms.Resize(transResize),  # 调整大小
            transforms.CenterCrop(transCrop),  # 中心裁剪
            transforms.ToTensor(),  # 转换为张量
            normalize  # 归一化
        ]
        transformSequence = transforms.Compose(transformList)

        # 创建测试集的数据加载器
        datasetTest = DatasetGenerator(pathImageDirectory=pathDirData,
                                       pathDatasetFile=pathFileTest,
                                       transform=transformSequence)
        dataLoaderTest = DataLoader(dataset=datasetTest, batch_size=trBatchSize,
                                    shuffle=False, num_workers=8,
                                    pin_memory=True)

        # 初始化张量来存储ground truth和预测结果
        outGT = torch.FloatTensor().cuda()
        outPRED = torch.FloatTensor().cuda()

        # 遍历测试集
        for i, (input, target) in enumerate(dataLoaderTest):
            target = target.cuda()
            input = input.cuda()

            with torch.no_grad():
                varInput = torch.autograd.Variable(input)
                out = model(varInput)  # 前向传播
                outPRED = torch.cat((outPRED, out), 0)  # 将输出结果连接起来
                outGT = torch.cat((outGT, target), 0)  # 将ground truth连接起来

        # 计算AUROC值
        aurocIndividual = ChexnetTrainer.computeAUROC(outGT, outPRED,
                                                      nnClassCount,
                                                      plot_roc_curve=True,
                                                      class_names=CLASS_NAMES)
        aurocMean = np.array(aurocIndividual).mean()  # 计算平均AUROC值

        # 输出每个类别的AUROC
        for i in range(len(aurocIndividual)):
            print(f'{CLASS_NAMES[i]}: {aurocIndividual[i]}')

        print(f'MEAN: {aurocMean}')