auxiliary_classifier.py

import os
#selected_gpus = [0] #configure this
#os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(gpu) for gpu in selected_gpus])

import torch
import math
import pandas as pd
import torchvision
from torchvision import transforms
import torch.nn as nn
from PIL import Image
from sklearn.metrics import mean_squared_error, mean_absolute_error
import argparse
import inversefed
from src.trainer import Trainer
from src.custom_models import DenseNet, ResNet
import wandb
from copy import deepcopy
from src.utils.model_utils import freeze_all_but_last, freeze_batchnorm, freeze_middle


DATA_BASE_PATH = '/dhc/dsets/ChestXrays/CheXpert/CheXpert-v1.0-large/CheXpert-v1.0'
#'/dhc/dsets/ChestXrays/CheXpert/'
image_size = 224
random_seed = 207


class CheXpertDataset(torch.utils.data.Dataset):
    # CheXpert mean and std
    xray_mean = 0.5029
    xray_std = 0.2899

    def __init__(self, folder_dir, dataframe, image_size, normalization):
        """
        Init Dataset

        Parameters
        ----------
        folder_dir: str
            folder contains all images
        dataframe: pandas.DataFrame
            dataframe contains all information of images
        image_size: int
            image size to rescale
        normalization: bool
            whether applying normalization with mean and std from ImageNet or not
        """
        self.image_paths = []  # List of image paths
        self.image_labels = []  # List of image labels

        # Define list of image transformations
        image_transformation = [
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor()
        ]

        if normalization:
            # Normalization with mean and std from ImageNet
            image_transformation.append(transforms.Normalize(self.xray_mean, self.xray_std))

        self.image_transformation = transforms.Compose(image_transformation)

        # Get all image paths and image labels from dataframe
        for _, row in dataframe.iterrows():
            self.image_paths.append(os.path.join(folder_dir, row['Path']))
            self.image_labels.append(row['Label'])

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        """
        Read image at index and convert to torch Tensor
        """

        # Read image
        image_path = self.image_paths[index]
        image_data = Image.open(image_path).convert('L')
        label = self.image_labels[index]

        # Resize and convert image to torch tensor
        image_data = self.image_transformation(image_data)

        return image_data, torch.FloatTensor([label])
        #return image_data, torch.tensor(label, dtype=torch.float)


class RMSELoss(torch.nn.modules.Module):
    __constants__ = ['reduction']

    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean') -> None:
        super().__init__()
        if size_average is not None or reduce is not None:
            self.reduction: str = torch.nn._reduction.legacy_get_string(size_average, reduce)
        else:
            self.reduction = reduction

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return torch.sqrt(torch.nn.functional.mse_loss(input, target, reduction=self.reduction))


def read_dataset(label, normalise=False):
    train_info = pd.read_csv(f'{DATA_BASE_PATH}/train.csv')
    valid_info = pd.read_csv(f'{DATA_BASE_PATH}/valid.csv')

    #train_info = pd.concat((train_info, valid_info))

    train_info = train_info[train_info['Frontal/Lateral'] == 'Frontal']
    valid_info = valid_info[valid_info['Frontal/Lateral'] == 'Frontal']

    train_info = train_info[['Path', label]].rename(columns={label: "Label"})
    valid_info = valid_info[['Path', label]].rename(columns={label: "Label"})

    if label == "Sex":
        train_info.replace({'Female': 0, 'Male': 1}, inplace=True)
        valid_info.replace({'Female': 0, 'Male': 1}, inplace=True)

        train_info = train_info[~(train_info['Label'] == 'Unknown')]
        valid_info = valid_info[~(valid_info['Label'] == 'Unknown')]
    elif normalise:
        max_age = 100
        print(f'Maximum Age: {max_age}')
        train_info.Label = train_info.Label.div(max_age)
        valid_info.Label = valid_info.Label.div(max_age)
    print('Example labels:')
    print(train_info.Label.head(5))

    train_dataset = CheXpertDataset(folder_dir=os.path.split(DATA_BASE_PATH)[0], dataframe=train_info, image_size=image_size, normalization=True)
    valid_dataset = CheXpertDataset(folder_dir=os.path.split(DATA_BASE_PATH)[0], dataframe=valid_info, image_size=image_size, normalization=True)

    return train_dataset, valid_dataset


# --------------------
# Evaluation metrics
# --------------------

# def compute_metrics(outputs, targets, losses):
#     n_classes = outputs.shape[1]
#     fpr, tpr, aucs, precision, recall = {}, {}, {}, {}, {}
#     for i in range(n_classes):
#         fpr[i], tpr[i], _ = roc_curve(targets[:,i], outputs[:,i])
#         aucs[i] = auc(fpr[i], tpr[i])
#         precision[i], recall[i], _ = precision_recall_curve(targets[:,i], outputs[:,i])
#         fpr[i], tpr[i], precision[i], recall[i] = fpr[i].tolist(), tpr[i].tolist(), precision[i].tolist(), recall[i].tolist()
#
#     metrics = {'fpr': fpr,
#                'tpr': tpr,
#                'aucs': aucs,
#                'precision': precision,
#                'recall': recall,
#                'loss': dict(enumerate(losses.mean(0).tolist()))}
#
#     return metrics


# --------------------
# Train and evaluate
# --------------------
@torch.no_grad()
def evaluate(model, dataloader, loss_fn, setup):
    model.eval()
    out_gt = torch.FloatTensor()
    out_pred = torch.FloatTensor()
    loss_val = 0
    for x, target in dataloader:
        out = model(x.to(**setup))
        loss = loss_fn(out, target.to(**setup))

        loss_val += loss.item()
        out_gt = torch.cat((out_gt, target), 0)
        out_pred = torch.cat((out_pred, out), 0)

    return out_gt, out_pred, loss_val / len(dataloader)


def evaluate_single_model(model, dataloader, loss_fn, setup):
    outputs, targets, loss = evaluate(model, dataloader, loss_fn, setup)
    outputs = outputs.cpu().numpy()
    targets = targets.cpu().numpy()
    if args.label == 'Sex':
        metrics = Trainer.compute_metrics(targets, outputs)
    else:
        metrics = {
            'MSE': mean_squared_error(targets, outputs),
            'MAE': mean_absolute_error(targets, outputs)
        }
    metrics['loss'] = loss
    return metrics
    #return compute_metrics(outputs, targets, losses)


def train_and_evaluate(model, train_dataloader, valid_dataloader, test_dataloader, loss_fn, optimizer, n_epochs, setup, freeze_mode='none'):
    if freeze_mode == 'batch_norm':
        freeze_batchnorm(model)
    if freeze_mode == 'all_but_last':
        freeze_all_but_last(model)
    if freeze_mode == 'middle':
        freeze_middle(model)

    for epoch in range(n_epochs):
        # train
        model.train()
        print(f'trainable vars: {sum(p.numel() for p in model.parameters() if p.requires_grad)}')
        #for x, target in train_dataloader:
        for i, (x, target) in enumerate(train_dataloader):
            out = model(x.to(**setup))
            loss = loss_fn(out, target.to(**setup))
            #predictions = out.argmax(dim=1, keepdim=True).squeeze()
            predictions = torch.bucketize(out, torch.tensor([0.5]), right=True)
            correct = (predictions.to(**setup) == target.to(**setup)).sum().item()
            accuracy = correct / args.batch_size

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # if scheduler and args.step >= args.lr_warmup_steps: scheduler.step()

            print(f'loss:{loss.item():.4f}, accuracy:{accuracy:.4f}')

            if i > 10:
                break
        print(f'X: {x.shape}')
        print(f'Target: {target.shape}')
        # evaluate
        print('Evaluating...', end='\r')
        eval_metrics = evaluate_single_model(model, valid_dataloader, loss_fn, setup)
        print(f'Evaluate metrics @ step {epoch}')
        print(f'AUC:{eval_metrics["aucs"]}\n')
        print(f'Loss:{eval_metrics["loss"]}\n')

        # save eval metrics
        #save_json(eval_metrics, 'eval_results_step_{}'.format(args.step), args)


def validate_regression_model(model, valid_dataloader, loss_fn, use_gpu, normalised):
    model.eval()
    lossVal = 0

    if use_gpu:
        out_gt = torch.FloatTensor().cuda()
        out_pred = torch.FloatTensor().cuda()
    else:
        out_gt = torch.FloatTensor()
        out_pred = torch.FloatTensor()

    with torch.no_grad():
        for model_input, target in valid_dataloader:

            if use_gpu:
                target = target.cuda(non_blocking=True)
                model_input = model_input.cuda(non_blocking=True)

            model_output = model(model_input)

            lossVal += loss_fn(model_output, target).item()

            # collect predictions and ground truth for AUROC computation
            out_gt = torch.cat((out_gt, target), 0)
            out_pred = torch.cat((out_pred, model_output), 0)

    # compute metrics
    out_gt_np = out_gt.cpu().numpy()
    out_pred_np = out_pred.cpu().numpy()
    metrics = {}

    if not normalised:
        metrics = {
            'Raw_MSE': mean_squared_error(out_gt_np, out_pred_np),
            'Raw_MAE': mean_absolute_error(out_gt_np, out_pred_np)
        }
        out_gt_np = out_gt_np / 100
        out_pred_np = out_pred_np / 100
    metrics = {
        **metrics,
        **{'MSE': mean_squared_error(out_gt_np, out_pred_np),
           'MAE': mean_absolute_error(out_gt_np, out_pred_np)}
    }
    if normalised:
        out_gt_np = out_gt_np * 100
        out_pred_np = out_pred_np * 100
        metrics = {
            **metrics,
            **{'Raw_MSE': mean_squared_error(out_gt_np, out_pred_np),
               'Raw_MAE': mean_absolute_error(out_gt_np, out_pred_np)}
        }

    return lossVal / len(valid_dataloader), metrics

def new_train_and_evaluate(model, train_dataloader, valid_dataloader, test_dataloader,
                           loss_fn, optimizer, n_epochs, es_epochs, freeze_mode, normalise_age):
    min_val_loss = 10000
    es_counter = 0
    es_model_state = None
    for epoch in range(n_epochs):
        losst = Trainer.epoch_train(model, train_dataloader, optimizer, loss_fn,
                                    use_gpu=torch.cuda.is_available(), freeze_mode=freeze_mode)
        print(f'Train Loss: {losst}')

        if args.label == 'Sex':
            lossv, metrics = Trainer.epoch_val(model, valid_dataloader, loss_fn, use_gpu=torch.cuda.is_available())
        else:
            lossv, metrics = validate_regression_model(model, valid_dataloader, loss_fn, use_gpu=torch.cuda.is_available(),normalised=normalise_age)
        print(f'Epoch {epoch} - Train Loss: {losst:.3f}, Val Loss: {lossv:.3f}')
        print(f'\tMetrics: {metrics}')
        wandb.log({
            **{f'val/{key}': val for key, val in metrics.items()},
            **{'train/train': losst,
               'val/loss': lossv}
        })

        # early stopping
        if lossv < min_val_loss:
            min_val_loss = lossv
            es_counter = 0
            es_model_state = deepcopy(model.state_dict())
        else:
            es_counter += 1
            if es_counter >= es_epochs:
                print(f'Early stopping in epoch {epoch} and resetting model')
                model.load_state_dict(es_model_state)
                break

    print('Final Evaluation')
    # Test model
    if args.label == 'Sex':
        loss_test, test_metrics = Trainer.epoch_val(model, test_dataloader, loss_fn, use_gpu=torch.cuda.is_available())
    else:
        loss_test, test_metrics = validate_regression_model(model, test_dataloader, loss_fn, use_gpu=torch.cuda.is_available(),
                                                            normalised=normalise_age)

    print(f'Eval Loss {loss_test}, Metrics: {test_metrics}')
    wandb.log({
        **{f'test/{key}': val for key, val in test_metrics.items()},
        **{'test/loss': loss_test}
    })

    return model


def get_argparser():
    parser = argparse.ArgumentParser(description='Predict either sex or age from CXR')

    parser.add_argument('--model', default='ResNet50', choices=['DenseNet121', 'ResNet50'], type=str,
                        help='Prediction model.')
    parser.add_argument('--trained_model', action='store_true', help='Use a trained model.')
    parser.add_argument('--save_model', action='store_true', help='Whether to save the trained model')
    parser.add_argument('--n_epochs', default=5, type=int, help='How many epochs to train?')
    parser.add_argument('--es_epochs', default=2, type=int, help='Num Epochs failure to increase loss for stopping early')
    parser.add_argument('--batch_size', default=64, type=int, help='Which batch size to use?')
    parser.add_argument('--optim', default='adam', type=str, help='Weigh the parameter list differently.')
    parser.add_argument('--lr', default=0.001, type=float, help='Learning rate')
    parser.add_argument('--betas', default=[0.9, 0.999], type=list, help='Adam Betas')
    parser.add_argument('--eps', default=1e-08, type=float, help='Adam Epsilon')
    parser.add_argument('--weight_decay', default=1e-5, type=float, help='Adam Weight Decay')
    parser.add_argument('--label', default='Sex', choices=['Sex', 'Age'], type=str, help='Which label to predict.')
    parser.add_argument('--normalise_age', type=str, default="False", help='Whether to normalise the age label')
    #parser.add_argument('--normalise_age', default=False, action='store_true', help='Whether to normalise the age label')
    parser.add_argument('--freeze', default='none', choices=['none', 'batch_norm', 'all_but_last'], type=str, help='Freeze Model Layers')
    parser.add_argument('--output_dir', default='logs', type=str, help='Where to print tensorboard log.')
    parser.add_argument('--deterministic', default=True)
    parser.add_argument('--wandb_disabled', default=True, action='store_false', dest='wandb')

    return parser


if __name__ == "__main__":
    parser = get_argparser()
    args = parser.parse_args()
    args.normalise_age = (args.normalise_age == "True")

    if args.wandb:
        wandb.init(project='CXR_Auxiliary_Classifier',
                   entity="bjarnepfitzner",
                   tags=[args.model, args.freeze, args.label], resume='allow',
                   config=args, allow_val_change=True,
                   settings=wandb.Settings(start_method="fork"))
    else:
        wandb.init(mode='disabled')

    setup = inversefed.utils.system_startup(args)

    # not entirely reproducible... only on GPU
    if args.deterministic:
        inversefed.utils.set_deterministic()
        inversefed.utils.set_random_seed(random_seed)

    num_classes = 1

    train_dataset, test_dataset = read_dataset(args.label, args.normalise_age)
    #complete_dataset = read_dataset(args.label, args.normalise_age)

    if args.deterministic:
        generator = torch.Generator().manual_seed(42)
    else:
        generator = torch.default_generator
    train_dataset, valid_dataset = torch.utils.data.random_split(train_dataset, [math.floor(0.9 * len(train_dataset)),
                                                                                 math.ceil(0.1 * len(train_dataset))],
                                                                 generator)
    print(f'Dataset sizes: {len(train_dataset)}, {len(valid_dataset)}, {len(test_dataset)}')

    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,
                                                   num_workers=8, pin_memory=True)
    valid_dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=100, shuffle=False,
                                                   num_workers=8, pin_memory=True)
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=100, shuffle=False,
                                                  num_workers=8, pin_memory=True)

    if args.model == 'DenseNet121':
        net = DenseNet
    elif args.model == 'ResNet50':
        net = ResNet
    else:
        exit('Model not supported')
    model = net(version=args.model, out_size=num_classes, input_layer_aggregation='mean', pre_trained=True,
                regression=(args.label=='Age')).cuda()

    optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr, betas=args.betas, weight_decay=args.weight_decay, eps=args.eps)

    loss_fn = torch.nn.BCELoss() if args.label == 'Sex' else RMSELoss()#torch.nn.RMSELoss()

    new_train_and_evaluate(model, train_dataloader, valid_dataloader, test_dataloader, loss_fn, optimizer,
                           args.n_epochs, args.es_epochs, args.freeze, args.normalise_age)

    if args.save_model:
        torch.save(model.state_dict(), f"models/AuxCls_{args.model}_{args.label}.pth.tar")