chenqitao/swinb

import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms, models from sklearn.metrics import confusion_matrix, roc_curve, auc, accuracy_score, ConfusionMatrixDisplay import os import numpy as np import matplotlib.pyplot as plt from torchvision.datasets import ImageFolder

数据路径

train_data_dir = "/root/autodl-tmp/vqvaeboi/train" val_data_dir = '/root/autodl-tmp/vqvaeboi/test'

best_accuracy = 0.0

数据预处理和加载

data_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), # 随机颜色抖动 transforms.RandomRotation(degrees=15), # 随机旋转 transforms.ToTensor(), # This should come before Normalize transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # This should be after ToTensor ])

train_dataset: ImageFolder = datasets.ImageFolder(root=train_data_dir, transform=data_transform) val_dataset = datasets.ImageFolder(root=val_data_dir, transform=data_transform)

dataloaders = { 'train': torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4), 'val': torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4), }

ResNet-18 模型定义

model = models.swin_b(pretrained=True) num_ftrs = model.head.in_features # 获取最后一个线性层输入的特征数量 model.head = nn.Linear(num_ftrs, 2) # 创建一个新的线性层，输出特征维度为 2

损失函数和优化器

criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.0001)

将模型和数据移动到GPU上

num_epochs = 100 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device)

train_losses = [] val_losses = [] train_accuracies = [] val_accuracies = []

best_val_labels = [] best_val_preds = []

for epoch in range(num_epochs): for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval()

    running_loss = 0.0
    running_corrects = 0
    all_labels = []
    all_preds = []

    for inputs, labels in dataloaders[phase]:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()

        with torch.set_grad_enabled(phase == 'train'):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            if phase == 'train':
                loss.backward()
                optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        _, preds = torch.max(outputs, 1)
        running_corrects += torch.sum(preds == labels.data)

        if phase == 'val':
            all_labels.extend(labels.cpu().numpy())
            all_preds.extend(preds.cpu().numpy())

    epoch_loss = running_loss / len(dataloaders[phase].dataset)
    epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

    print(f'Epoch {epoch}/{num_epochs} | {phase} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f}')

    if phase == 'train':
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)
    else:
        val_losses.append(epoch_loss)
        val_accuracies.append(epoch_acc)
        if epoch_acc > best_accuracy:
            best_accuracy = epoch_acc
            torch.save(model.state_dict(), 'best_model.pth')
            best_val_labels = all_labels
            best_val_preds = all_preds

加载最佳模型参数

model.load_state_dict(torch.load('best_model.pth'))

可视化部分...

绘制 Loss 曲线图

plt.figure(figsize=(10, 5)) plt.plot(train_losses, label='Training Loss') plt.plot(val_losses, label='Validation Loss') plt.xlabel('Epochs') plt.ylabel('Loss') plt.title('Training and Validation Loss') plt.legend() plt.show()

绘制 Accuracy 曲线图

plt.figure(figsize=(10, 5)) plt.plot(train_accuracies, label='Training Accuracy') plt.plot(val_accuracies, label='Validation Accuracy') plt.xlabel('Epochs') plt.ylabel('Accuracy') plt.title('Training and Validation Accuracy') plt.legend() plt.show()

使用 best_val_labels 和 best_val_preds 来计算混淆矩阵

conf_matrix = confusion_matrix(best_val_labels, best_val_preds) print("Confusion Matrix:") print(conf_matrix)

绘制混淆矩阵图

class_names = sorted(train_dataset.classes) disp = ConfusionMatrixDisplay(confusion_matrix=conf_matrix, display_labels=class_names) disp.plot(cmap='Blues', values_format='d') plt.title('Confusion Matrix') plt.show()

计算 ROC 曲线和 AUC

fpr, tpr, _ = roc_curve(best_val_labels, best_val_preds) roc_auc = auc(fpr, tpr)

首先计算混淆矩阵

conf_matrix = confusion_matrix(best_val_labels, best_val_preds)

对于二分类，混淆矩阵布局如下：

TN FP

FN TP

TN, FP, FN, TP = conf_matrix.ravel()

计算特异度和灵敏度

Specificity = TN / (TN + FP) Sensitivity = TP / (TP + FN)

print(f'Specificity: {Specificity:.4f}') print(f'Sensitivity: {Sensitivity:.4f}')

绘制 ROC 曲线

plt.figure(figsize=(8, 8)) plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'AUC = {roc_auc:.2f}') plt.plot([0, 1], [0, 1], color='navy', lw=2, 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') plt.legend(loc="lower right") plt.show()

输出最佳准确度

print(f'Best Validation Accuracy: {best_accuracy:.4f}')