import os import numpy as np import torch import torch.nn as nn import torch.optim as optim from torchvision import transforms import matplotlib.pyplot as plt from torch.utils.data import DataLoader from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score, f1_score from tqdm import tqdm import joblib # 导入joblib库用于保存模型 from DatasetGenerator import DatasetGenerator from DensenetModels import DenseNet121, DenseNet169, DenseNet201, ResNet50, FocalLoss from sklearn.multioutput import MultiOutputClassifier from sklearn.metrics import roc_curve, auc class ChexnetTrainer: # 在 ChexnetTrainer 类中添加以下方法 @staticmethod def train(pathDirData, pathFileTrain, pathFileVal, nnArchitecture, nnIsTrained, nnClassCount, trBatchSize, trMaxEpoch, transResize, transCrop, launchTimestamp, checkpoint=None): """训练函数""" # 自动选择设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 获取模型并迁移到设备 model = ChexnetTrainer._get_model(nnArchitecture, nnClassCount, nnIsTrained).to(device) # 设置数据预处理和数据加载器 transformSequence = ChexnetTrainer._get_transform(transCrop) datasetTrain = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTrain, transform=transformSequence, model=model) datasetVal = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileVal, transform=transformSequence, model=model) dataLoaderTrain = DataLoader(dataset=datasetTrain, batch_size=trBatchSize, shuffle=True, num_workers=4, pin_memory=True) dataLoaderVal = DataLoader(dataset=datasetVal, batch_size=trBatchSize, shuffle=False, num_workers=4, pin_memory=True) optimizer = optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5) # 替换 ReduceLROnPlateau 为 StepLR scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # 迁移class_weights到对应设备 class_weights = torch.tensor( [1.1762, 0.6735, 0.9410, 1.6680, 0.9699, 1.1950, 2.2584, 0.6859, 1.6683, 0.7744, 0.4625, 0.7385, 1.3764, 0.1758], dtype=torch.float32).to(device) loss = FocalLoss(alpha=class_weights, gamma=2, logits=True) lossMIN = 100000 train_f1_scores, val_f1_scores = [], [] all_val_targets = [] all_val_scores = [] # 训练循环 for epochID in range(trMaxEpoch): ChexnetTrainer.epochTrain(model, dataLoaderTrain, optimizer, scheduler, trMaxEpoch, nnClassCount, loss, train_f1_scores, device) lossVal, losstensor, val_f1, val_targets, val_scores = ChexnetTrainer.epochVal( model, dataLoaderVal, optimizer, scheduler, trMaxEpoch, nnClassCount, loss, val_f1_scores, device) all_val_targets.append(val_targets) all_val_scores.append(val_scores) # 更新学习率 scheduler.step() # 保存最佳模型 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') # 合并所有验证集的标签和预测 all_val_targets = np.vstack(all_val_targets) all_val_scores = np.vstack(all_val_scores) # 绘制F1分数图 plt.figure() 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(os.path.join('images', 'f1_scores.png')) plt.close() # 计算并绘制AUC-ROC曲线 CLASS_NAMES = ['Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia'] ChexnetTrainer.plot_auc_roc(all_val_targets, all_val_scores, CLASS_NAMES, os.path.join('images', 'auc_roc_curve.png')) print("AUC-ROC曲线已保存到 images/auc_roc_curve.png") # 提取特征并训练随机森林 rf_classifier = ChexnetTrainer.train_random_forest(datasetTrain, datasetVal, model, device) # 保存随机森林模型 joblib.dump(rf_classifier, os.path.join('images', 'random_forest_model.pkl')) # 保存随机森林分类器 @staticmethod def epochTrain(model, dataLoader, optimizer, scheduler, epochMax, classCount, loss, f1_scores, device): model.train() all_targets = [] all_preds = [] # 创建进度条 pbar = tqdm(total=len(dataLoader), desc="Training", leave=True) for batchID, (input, target) in enumerate(dataLoader): target = target.to(device) input = input.to(device) 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) # 更新进度条 pbar.update(1) pbar.close() f1 = f1_score(np.array(all_targets), np.array(all_preds), average="macro") f1_scores.append(f1) # 在每个epoch结束时打印当前的F1-score print(f"Epoch completed. F1 Score (Macro): {f1:.4f}") @staticmethod def epochVal(model, dataLoader, optimizer, scheduler, epochMax, classCount, loss, f1_scores, device): model.eval() lossVal = 0 lossValNorm = 0 losstensorMean = 0 all_targets = [] all_preds = [] all_scores = [] # 收集预测概率 # 创建进度条 pbar = tqdm(total=len(dataLoader), desc="Validation", leave=True) for i, (input, target) in enumerate(dataLoader): target = target.to(device) input = input.to(device) 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 scores = torch.sigmoid(varOutput).cpu().data.numpy() all_scores.extend(scores) pred = scores > 0.5 all_targets.extend(target.cpu().numpy()) all_preds.extend(pred) # 更新进度条 pbar.update(1) pbar.close() f1 = f1_score(np.array(all_targets), np.array(all_preds), average="macro") f1_scores.append(f1) # 在每个epoch结束时打印当前的F1-score print(f"Epoch completed. F1 Score (Macro): {f1:.4f}") return lossVal / lossValNorm, losstensorMean / lossValNorm, f1, np.array(all_targets), np.array(all_scores) @staticmethod def train_random_forest(datasetTrain, datasetVal, model, device): """训练随机森林分类器，适应多标签分类""" print("正在提取训练数据的特征。") # 提取训练数据的特征 train_features, train_labels = ChexnetTrainer.extract_features(datasetTrain, model, device) val_features, val_labels = ChexnetTrainer.extract_features(datasetVal, model, device) print("正在训练随机森林分类器。") # 使用MultiOutputClassifier来处理多标签问题 rf_classifier = MultiOutputClassifier(RandomForestClassifier(n_estimators=100)) rf_classifier.fit(train_features, train_labels) print("正在评估随机森林分类器。") # 在验证集上评估随机森林 val_preds = rf_classifier.predict(val_features) val_f1 = f1_score(val_labels, val_preds, average='macro') print(f"Random Forest F1 Score on validation set: {val_f1}") return rf_classifier @staticmethod def extract_features(dataset, model, device): """提取数据集的特征""" features = [] labels = [] print("正在提取数据集的特征。") model.eval() dataLoader = DataLoader(dataset=dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True) for input, target in tqdm(dataLoader, desc="Extracting Features"): input = input.to(device) target = target.to(device) with torch.no_grad(): varInput = torch.autograd.Variable(input) varOutput = model(varInput) # 这里假设 varOutput 是一个二维张量，我们需要把它展平成一维向量 features.append(varOutput.view(varOutput.size(0), -1).cpu().data.numpy()) # 展平特征 labels.append(target.cpu().data.numpy()) # 假设 target 是多标签格式的 # 使用 np.vstack 将特征和标签拼接成数组 return np.vstack(features), np.vstack(labels) # 确保标签是二维矩阵，每一行对应一个样本的标签 @staticmethod def plot_auc_roc(y_true, y_scores, class_names, save_path): """绘制AUC-ROC曲线并保存""" plt.figure(figsize=(20, 15)) for i, class_name in enumerate(class_names): fpr, tpr, _ = roc_curve(y_true[:, i], y_scores[:, i]) roc_auc = auc(fpr, tpr) plt.plot(fpr, tpr, lw=2, label=f'ROC curve of class {class_name} (area = {roc_auc:.2f})') plt.plot([0, 1], [0, 1], 'k--', lw=2) 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(save_path) plt.close() @staticmethod def test(pathDirData, pathFileTest, pathModel, pathRfModel, nnArchitecture, nnClassCount, nnIsTrained, trBatchSize, imgtransResize, imgtransCrop, timestampLaunch): """测试函数，支持深度学习模型和随机森林模型""" # 加载深度学习模型 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ChexnetTrainer._get_model(nnArchitecture, nnClassCount, nnIsTrained).to(device) checkpoint = torch.load(pathModel, map_location=device) model.load_state_dict(checkpoint['state_dict']) model.eval() # 确保模型处于评估模式 # 加载随机森林模型 rf_classifier = joblib.load(pathRfModel) # 获取 transform transformSequence = ChexnetTrainer._get_transform(imgtransCrop) datasetTest = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTest, transform=transformSequence, model=model) # 使用 extract_features 提取特征和标签 features, labels = ChexnetTrainer.extract_features(datasetTest, model, device) # 使用随机森林模型进行预测 rf_preds = rf_classifier.predict(features) # 计算多标签分类的 F1 分数 rf_f1 = f1_score(labels, rf_preds, average="macro") print(f"Random Forest Multi-label F1 Score on test set: {rf_f1}") # 计算AUC-ROC rf_scores = [] if hasattr(rf_classifier, "predict_proba"): rf_scores = rf_classifier.predict_proba(features) # 将列表转换为数组，并选择正类的概率 rf_scores = np.array([rf_scores[i][:, 1] for i in range(len(rf_scores))]).reshape(labels.shape) else: # 如果没有predict_proba方法，则无法计算AUC print("随机森林模型不支持predict_proba方法，无法计算AUC-ROC。") rf_scores = np.zeros_like(labels) # 绘制AUC-ROC曲线 CLASS_NAMES = ['Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia'] ChexnetTrainer.plot_auc_roc(labels, rf_scores, CLASS_NAMES, os.path.join('images', 'test_auc_roc_curve.png')) print("测试集的AUC-ROC曲线已保存到 images/test_auc_roc_curve.png") # 输出深度学习模型和随机森林模型的 F1 分数 print(f"Testing completed. Random Forest F1 Score: {rf_f1}") return labels, rf_preds # 返回真实标签和预测结果 # 其他方法保持不变... @staticmethod def _get_model(nnArchitecture, nnClassCount, nnIsTrained): """根据选择的模型架构返回对应的模型""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 根据模型架构选择模型 if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).to(device) elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).to(device) elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).to(device) elif nnArchitecture == 'RESNET-50': model = ResNet50(nnClassCount, nnIsTrained).to(device) else: raise ValueError( f"Unknown architecture: {nnArchitecture}. Please choose from 'DENSE-NET-121', 'DENSE-NET-169', 'DENSE-NET-201', 'RESNET-50'.") return model @staticmethod def _get_transform(transCrop): """返回图像预处理的转换""" normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) transformList = [ transforms.RandomResizedCrop(transCrop), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize ] return transforms.Compose(transformList)