ADD file via upload

Resnet_SEBlock_Focalloss.py

@@ -0,0 +1,589 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms
+from torchvision.models import resnet34
+from torch.utils.data import random_split
+from torch.utils.data import WeightedRandomSampler
+from sklearn.metrics import precision_score, recall_score, f1_score
+from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
+from tqdm import tqdm
+import time
+import random
+import os
+import shutil
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import pprint
+import timm
+import pytorch_lightning as pl
+import logging
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import seaborn as sns
+
+h = {
+    "num_epochs": 6,
+    "batch_size": 16,
+    "image_size": 224,
+    "fc1_size": 512,
+    "lr": 0.001,
+    "model": "resnet34_custom",
+    "scheduler": "CosineAnnealingLR10",
+    "balance": True,
+    "early_stopping_patience": float("inf"),
+    "use_best_checkpoint": False
+}
+
+class CustomImageFolder(torch.utils.data.Dataset):
+    def __init__(self, root, transform=None, is_valid_file=None):
+        self.dataset = datasets.ImageFolder(root, is_valid_file=is_valid_file)
+        self.transform = transform
+        self.targets = self.dataset.targets
+
+    def __getitem__(self, index):
+        image, label = self.dataset[index]
+        if self.transform:
+            image = self.transform(image=np.array(image))["image"] / 255.0
+        return image, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+class PneumoniaDataModule(pl.LightningDataModule):
+    def __init__(self, h, data_dir):
+        super().__init__()
+        self.h = h
+        self.data_dir = data_dir
+
+    def setup(self, stage=None):
+        data_transforms_train_alb = A.Compose([
+            A.Rotate(limit=20),
+            A.HorizontalFlip(p=0.5),
+            A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=1),
+            A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0, rotate_limit=0, p=0.5),
+            A.Perspective(scale=(0.05, 0.15), keep_size=True, p=0.5),
+            A.Resize(height=h["image_size"], width=h["image_size"]),
+            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ToTensorV2()
+        ])        
+
+        data_transforms_val_alb = A.Compose([
+            A.Resize(self.h["image_size"], self.h["image_size"]),
+            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ToTensorV2(),
+        ])  
+
+        val_split = 0.2
+        train_filenames, val_filenames = self._split_file_names(self.data_dir+"train/", val_split)
+
+        # Load the datasets
+        self.train_dataset = CustomImageFolder(self.data_dir+"train/", transform=data_transforms_train_alb, is_valid_file=lambda x: x in train_filenames)
+        self.val_dataset = CustomImageFolder(self.data_dir+"train/", transform=data_transforms_val_alb, is_valid_file=lambda x: x in val_filenames)    
+        self.test_dataset = CustomImageFolder(self.data_dir+"test/", transform=data_transforms_val_alb, is_valid_file=lambda x: self._is_image_file(x))
+    
+    def train_dataloader(self):
+        if self.h["balance"]:
+            sampler = self._create_weighted_sampler(self.train_dataset)
+            return torch.utils.data.DataLoader(self.train_dataset, batch_size=self.h["batch_size"], sampler=sampler, num_workers=0)
+        else:
+            return torch.utils.data.DataLoader(self.train_dataset, batch_size=self.h["batch_size"], shuffle=True, num_workers=0)
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(self.val_dataset, batch_size=self.h["batch_size"], num_workers=0)
+
+    def test_dataloader(self):
+        return torch.utils.data.DataLoader(self.test_dataset, batch_size=self.h["batch_size"], num_workers=0)
+
+    def _extract_patient_ids(self, filename):
+        patient_id = filename.split('_')[0].replace("person", "")
+        return patient_id
+
+    def _is_image_file(self, file_path):
+        return file_path.lower().endswith((".jpeg", ".jpg", ".png"))
+
+    def _split_file_names(self, input_folder, val_split_perc):
+        # Pneumonia files contain patient id, so we group split them by patient to avoid data leakage
+        pneumonia_patient_ids = set([self._extract_patient_ids(fn) for fn in os.listdir(os.path.join(input_folder, 'PNEUMONIA'))])
+        pneumonia_val_patient_ids = random.sample(list(pneumonia_patient_ids), int(val_split_perc * len(pneumonia_patient_ids)))
+
+        pneumonia_val_filenames = []
+        pneumonia_train_filenames = []
+
+        for filename in os.listdir(os.path.join(input_folder, 'PNEUMONIA')):
+            if self._is_image_file(filename):
+                patient_id = self._extract_patient_ids(filename)
+                if patient_id in pneumonia_val_patient_ids:
+                    pneumonia_val_filenames.append(os.path.join(input_folder, 'PNEUMONIA', filename))
+                else:
+                    pneumonia_train_filenames.append(os.path.join(input_folder, 'PNEUMONIA', filename))
+
+        # Normal (by file, no patient information in file names)
+        normal_filenames  = [os.path.join(input_folder, 'NORMAL', fn) for fn in os.listdir(os.path.join(input_folder, 'NORMAL'))]
+        normal_filenames = [filename for filename in normal_filenames if self._is_image_file(filename)]
+        normal_val_filenames = random.sample(normal_filenames, int(val_split_perc * len(normal_filenames)))
+        normal_train_filenames = list(set(normal_filenames)-set(normal_val_filenames))
+
+        train_filenames = pneumonia_train_filenames + normal_train_filenames
+        val_filenames = pneumonia_val_filenames + normal_val_filenames
+
+        return train_filenames, val_filenames        
+
+
+    def _create_weighted_sampler(self, dataset):
+        targets = dataset.targets
+        class_counts = np.bincount(targets)
+        class_weights = 1.0 / class_counts
+        weights = [class_weights[label] for label in targets]
+        sampler = WeightedRandomSampler(weights, len(weights))
+        return sampler
+
+#带seblock的Resnet34
+class SEBlock(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SEBlock, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+    
+class SEBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1, downsample=None):
+        super(SEBasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.se = SEBlock(planes)
+        self.downsample = downsample
+
+    def forward(self, x):
+        identity = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out = self.se(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        out = self.relu(out)
+        return out
+
+def make_layer(block, in_planes, planes, blocks, stride=1):
+    downsample = None
+    if stride != 1 or in_planes != planes * block.expansion:
+        downsample = nn.Sequential(
+            nn.Conv2d(in_planes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
+            nn.BatchNorm2d(planes * block.expansion),
+        )
+
+    layers = [block(in_planes, planes, stride, downsample)]
+    for _ in range(1, blocks):
+        layers.append(block(planes * block.expansion, planes))
+
+    return nn.Sequential(*layers)
+
+class CustomResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=2):
+        super(CustomResNet, self).__init__()
+        self.in_planes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = make_layer(block, 64, 64, layers[0])
+        self.layer2 = make_layer(block, 64, 128, layers[1], stride=2)
+        self.layer3 = make_layer(block, 128, 256, layers[2], stride=2)
+        self.layer4 = make_layer(block, 256, 512, layers[3], stride=2)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+        return x
+
+#替换损失函数
+class FocalLoss(nn.Module):
+    def __init__(self, alpha=1, gamma=2, reduction='mean'):
+        super(FocalLoss, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, inputs, targets):
+        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
+        pt = torch.exp(-ce_loss)
+        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
+        return focal_loss.mean() if self.reduction == 'mean' else focal_loss.sum()
+
+class PneumoniaModel(pl.LightningModule):
+    def __init__(self, h):
+        super().__init__()
+        self.h = h
+        self.model = self._create_model()
+        self.criterion = FocalLoss(alpha=1, gamma=2)
+        self.test_outputs = []
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(inputs)
+        loss = self.criterion(outputs, labels)
+        self.log('train_loss', loss, on_epoch=True, on_step=True, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(inputs)
+        loss = self.criterion(outputs, labels)
+        acc = (outputs.argmax(dim=1) == labels).float().mean()
+        metrics = {"val_loss": loss, "val_acc": acc}
+        self.log_dict(metrics, on_epoch=True, on_step=True, prog_bar=True)
+        return metrics        
+
+    def on_test_epoch_start(self):
+        self.test_outputs = []
+
+    def test_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(inputs)
+        loss = self.criterion(outputs, labels)
+        acc = (outputs.argmax(dim=1) == labels).float().mean()
+        preds = torch.argmax(outputs, dim=1)    
+        self.test_outputs.append({"test_loss": loss, "test_acc": acc, "preds": preds, "labels": labels})
+        return {"test_loss": loss, "test_acc": acc, "preds": preds, "labels": labels}
+
+    def on_test_epoch_end(self):
+        test_loss_mean = torch.stack([x["test_loss"] for x in self.test_outputs]).mean()
+        test_acc_mean = torch.stack([x["test_acc"] for x in self.test_outputs]).mean()
+
+        self.test_predicted_labels = torch.cat([x["preds"] for x in self.test_outputs], dim=0).cpu().numpy()
+        self.test_true_labels = torch.cat([x["labels"] for x in self.test_outputs], dim=0).cpu().numpy()
+        #Todo: remove f1 calculation from here
+        f1 = f1_score(self.test_true_labels, self.test_predicted_labels)
+
+        self.test_f1 = f1
+        self.test_acc = test_acc_mean.cpu().numpy() #Todo - fix it
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.h["lr"])
+        scheduler_dic = self._configure_scheduler(optimizer)
+
+        if (scheduler_dic["scheduler"]):
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": scheduler_dic
+            }            
+        else:
+            return optimizer
+
+    def _configure_scheduler(self, optimizer):
+        scheduler_name = self.h["scheduler"]
+        lr = self.h["lr"]
+        if (scheduler_name==""):
+            return {
+                "scheduler": None
+            }
+        if (scheduler_name=="CosineAnnealingLR10"):
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=h["num_epochs"], eta_min=lr*0.1) #*len(train_loader) if "step"
+            return {
+                "scheduler": scheduler,
+                "interval": "epoch"
+            }
+        if (scheduler_name=="ReduceLROnPlateau5"):
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, verbose=True)
+            return {
+                "scheduler": scheduler,
+                "interval": "epoch",
+                "monitor": "val_loss",
+                "strict": True
+            }
+        print ("Error. Unknown scheduler name '{scheduler_name}'")
+        return None
+
+    def _create_model(self):
+        if (self.h["model"]=="efficientnetv2"):
+            return timm.create_model("tf_efficientnetv2_b0", pretrained=True, num_classes=2)
+        if (self.h["model"]=="fc"):
+            return nn.Sequential(
+                nn.Flatten(),
+                nn.Linear(3 * self.h["image_size"] * self.h["image_size"], self.h["fc1_size"]),
+                nn.ReLU(),
+                nn.Linear(self.h["fc1_size"], 2)
+            )
+        if (self.h["model"]=="cnn"):
+            return nn.Sequential(
+                nn.Conv2d(3, 16, 3, padding=1),
+                nn.ReLU(),
+                nn.MaxPool2d(2, 2),
+                nn.Conv2d(16, 32, 3, padding=1),
+                nn.ReLU(),
+                nn.MaxPool2d(2, 2),
+                nn.Conv2d(32, 64, 3, padding=1),
+                nn.ReLU(),
+                nn.MaxPool2d(2, 2),
+                nn.Flatten(),
+                nn.Dropout(0.25),
+                nn.Linear(64 * (self.h["image_size"] // 8) * (self.h["image_size"] // 8), 512),
+                nn.ReLU(),
+                nn.Dropout(0.25),
+                nn.Linear(512, 2)
+            )
+        if (self.h["model"]=="resnet34_custom"):
+        
+            model = CustomResNet(SEBasicBlock, [3, 4, 6, 3], num_classes=2)
+            return model     
+
+class InfoPrinterCallback(pl.Callback):
+    def __init__(self):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if torch.cuda.is_available():
+            print(f"CPU cores: {os.cpu_count()}, Device: {device}, GPU: {torch.cuda.get_device_name(0)}")
+        else:
+            print(f"CPU cores: {os.cpu_count()}, Device: {device}")
+
+        # Print hyperparameters for records
+        print("Hyperparameters:")
+        pprint.pprint(h, indent=4)
+
+    def setup(self, trainer, pl_module, stage):
+        self.start_time = time.time()     
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        # Skip the sanity check
+        if trainer.sanity_checking:
+            return
+
+        epoch = trainer.current_epoch
+        total_epochs = trainer.max_epochs
+
+        elapsed_time = time.time() - self.start_time
+        avg_time_per_epoch = elapsed_time / (epoch + 1)
+        avg_time_per_epoch_min, avg_time_per_epoch_sec = divmod(avg_time_per_epoch, 60)
+
+        remaining_epochs = total_epochs - epoch - 1
+        remaining_time = remaining_epochs * avg_time_per_epoch
+        remaining_time_min, remaining_time_sec = divmod(remaining_time, 60)
+
+        print(f"Epoch {epoch + 1}/{total_epochs}: ", end="")
+
+        if "val_loss" in trainer.callback_metrics:
+            validation_loss = trainer.callback_metrics["val_loss"].cpu().numpy()
+            #self.validation_losses.append(validation_loss)            
+            print(f"Validation Loss = {validation_loss:.4f}", end="")
+        else:
+            print(f"Validation Loss not available", end="")
+
+        if "train_loss_epoch" in trainer.logged_metrics:
+            train_loss = trainer.logged_metrics["train_loss_epoch"].cpu().numpy()
+            print(f", Train Loss = {train_loss:.4f}", end="")
+        else:
+            print(f", Train Loss not available", end="")
+
+        print(f", Epoch Time: {avg_time_per_epoch_min:.0f}m {avg_time_per_epoch_sec:02.0f}s, Remaining Time: {remaining_time_min:.0f}m {remaining_time_sec:02.0f}s")
+
+    def plot_losses(self):
+        plt.style.use("seaborn-v0_8-whitegrid")
+        plt.rcParams.update({
+            "font.family": "serif",
+            "font.size": 10,
+            "axes.titlesize": 11,
+            "axes.labelsize": 10,
+            "xtick.labelsize": 9,
+            "ytick.labelsize": 9,
+            "legend.fontsize": 9,
+            "figure.figsize": (4, 3),
+            "figure.dpi": 150,
+            "lines.linewidth": 2,
+            "lines.markersize": 5
+        })
+
+        plt.figure()
+        plt.plot(self.training_losses, label="Training Loss", marker='o')
+        plt.plot(self.validation_losses, label="Validation Loss", marker='s')
+        plt.xlabel("Epoch")
+        plt.ylabel("Loss")
+        plt.legend()
+        plt.tight_layout()
+        plt.show() 
+
+class PlotTestConfusionMatrixCallback(pl.Callback):
+    def on_test_end(self, trainer, pl_module):
+        cm = confusion_matrix(pl_module.test_true_labels, pl_module.test_predicted_labels)
+        class_names = ["Normal", "Pneumonia"]
+
+        plt.style.use("seaborn-v0_8-white")
+        plt.rcParams.update({
+            "font.family": "serif",
+            "font.size": 10,
+            "axes.titlesize": 11,
+            "axes.labelsize": 10,
+            "xtick.labelsize": 9,
+            "ytick.labelsize": 9,
+            "legend.fontsize": 9,
+            "figure.figsize": (4, 3),
+            "figure.dpi": 150
+        })
+
+        fig, ax = plt.subplots()
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
+                    xticklabels=class_names,
+                    yticklabels=class_names,
+                    cbar=False, linewidths=0.5, linecolor='gray')
+        ax.set_xlabel("Predicted Label")
+        ax.set_ylabel("True Label")
+        ax.set_title("Confusion Matrix")
+        plt.tight_layout()
+        plt.show()
+
+class PlotTrainingLogsCallback(pl.Callback):
+    def __init__(self):
+        self.validation_losses = []
+        self.training_losses = []
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        if "train_loss_epoch" in trainer.logged_metrics:
+            train_loss = trainer.logged_metrics["train_loss_epoch"].cpu().numpy()
+            self.training_losses.append(train_loss)
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if trainer.sanity_checking:
+            return
+
+        if "val_loss" in trainer.callback_metrics:
+            validation_loss = trainer.callback_metrics["val_loss"].cpu().numpy()
+            self.validation_losses.append(validation_loss)            
+
+    def on_fit_end(self, trainer, pl_module):
+        plt.style.use("seaborn-v0_8-whitegrid")
+        plt.rcParams.update({
+            "font.family": "serif",
+            "font.size": 10,
+            "axes.titlesize": 11,
+            "axes.labelsize": 10,
+            "xtick.labelsize": 9,
+            "ytick.labelsize": 9,
+            "legend.fontsize": 9,
+            "figure.figsize": (4, 3),
+            "figure.dpi": 150,
+            "lines.linewidth": 2,
+            "lines.markersize": 5
+        })
+
+        plt.figure()
+        plt.plot(self.training_losses, label="Training Loss", marker='o')
+        plt.plot(self.validation_losses, label="Validation Loss", marker='s')
+        plt.xlabel("Epoch")
+        plt.ylabel("Loss")
+        plt.legend()
+        plt.tight_layout()
+        plt.show()  
+
+def check_solution(h, verbose):
+    pneumonia_data = PneumoniaDataModule(h, "D:/DATASET/archive/chest_xray/chest_xray/")
+    pneumonia_model = PneumoniaModel(h)
+
+    # Callbacks
+    info_printer = InfoPrinterCallback()
+
+    early_stopping = pl.callbacks.EarlyStopping(
+        monitor="val_loss",
+        patience=h["early_stopping_patience"],
+        verbose=True,
+    )
+
+    checkpoint_callback = pl.callbacks.ModelCheckpoint(
+        dirpath="model_checkpoints",
+        monitor="val_loss",
+        verbose=True,
+    )
+
+    callbacks = [info_printer, early_stopping, checkpoint_callback]
+    if (verbose):
+        callbacks.append(PlotTestConfusionMatrixCallback())
+        callbacks.append(PlotTrainingLogsCallback())
+
+    trainer = pl.Trainer(
+        max_epochs=h["num_epochs"],
+        accelerator="auto",
+        callbacks=callbacks,
+        log_every_n_steps=1,
+        fast_dev_run=False
+    )
+
+    trainer.fit(pneumonia_model, datamodule=pneumonia_data)
+
+    if (h["use_best_checkpoint"]):
+        #Debug lines
+        trainer.test(pneumonia_model, datamodule=pneumonia_data)
+        print(f"Last: F1= {pneumonia_model.test_f1:.4f}, Acc= {pneumonia_model.test_acc:.4f}")
+
+        best_model_path = checkpoint_callback.best_model_path
+        best_model = PneumoniaModel.load_from_checkpoint(best_model_path, h=h)
+        pneumonia_model = best_model
+
+    trainer.test(pneumonia_model, datamodule=pneumonia_data)
+    print(f"Best: F1= {pneumonia_model.test_f1:.4f}, Acc= {pneumonia_model.test_acc:.4f}")
+
+
+    return pneumonia_model.test_f1, pneumonia_model.test_acc
+
+
+logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
+
+f1_array = np.array([])
+accuracy_array = np.array([])
+start_time = time.time()
+
+repeats = 1
+for i in range(repeats):
+    print("===============================================")
+    print(f"Running solution {i+1}/{repeats}")
+    f1, accuracy = check_solution(h, verbose=(i==0))
+    print(f"F1 = {f1:.2f}, accuracy = {accuracy:.2f} ")
+    f1_array = np.append(f1_array, f1)
+    accuracy_array = np.append(accuracy_array, accuracy) 
+
+# Calculate elapsed time and remaining time
+repeat_time = (time.time() - start_time) / repeats
+repeat_time_min, repeat_time_sec = divmod(repeat_time, 60)
+
+# Printing final results
+print("Results")
+print(f"F1: {np.mean(f1_array):.1%} (+-{np.std(f1_array):.1%})")
+print(f"Accuracy: {np.mean(accuracy_array):.1%} (+-{np.std(accuracy_array):.1%})")
+print(f"Time of one solution: {repeat_time_min:.0f}m {repeat_time_sec:.0f}s")
+print(f" | {np.mean(f1_array):.1%} (+-{np.std(f1_array):.1%}) | {np.mean(accuracy_array):.1%} (+-{np.std(accuracy_array):.1%}) | {repeat_time_min:.0f}m {repeat_time_sec:.0f}s")
+
+# Print hyperparameters for reminding what the final data is for
+print("Hyperparameters:")
+pprint.pprint(h, indent=4)