train.py

import os
from typing import (Any, List, Dict, Optional, Tuple,
                    Union, Callable, Iterable, Iterator)
import pandas as pd
from PIL import Image
import datetime
from argparse import ArgumentParser
from enum import Enum
import numpy as np
from numpy.random import RandomState
import collections.abc
from collections import Counter, defaultdict
import math

import torch
import torch.nn as nn
import torch.utils.data as data
from torch.utils.data import DataLoader

from torchvision.transforms import (
    CenterCrop, 
    Compose, 
    Normalize, 
    RandomHorizontalFlip,
    RandomResizedCrop, 
    RandomRotation,
    RandomAffine,
    Resize, 
    ToTensor)

from transformers import ViTImageProcessor
from transformers import ViTForImageClassification
from transformers import AdamW

from transformers import AutoImageProcessor, ResNetForImageClassification

import lightning as L
from lightning import Trainer
from lightning.pytorch.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import ModelSummary
from torchmetrics.aggregation import MeanMetric
from torchmetrics.classification.accuracy import MulticlassAccuracy
from torchmetrics.classification import MulticlassCohenKappa

from labelmap import DR_LABELMAP


DataRecord = Tuple[Image.Image, int]


class RetinopathyDataset(data.Dataset[DataRecord]):
    def __init__(self, data_path: str) -> None:
        super().__init__()

        self.data_path = data_path

        self.ext = ".jpeg"

        anno_path = os.path.join(data_path, "trainLabels.csv")
        self.anno_df = pd.read_csv(anno_path) # ['image', 'level']
        anno_name_set = set(self.anno_df['image']) 

        if True:
            train_path = os.path.join(data_path, "train")
            img_path_list = os.listdir(train_path)
            img_name_set = set([os.path.splitext(p)[0] for p in img_path_list])
            assert anno_name_set == img_name_set

        self.label_map = DR_LABELMAP
    
    def __getitem__(self, index: Union[int, slice]) -> DataRecord:
        assert isinstance(index, int)
        img_path = self.get_path_at(index)
        img = Image.open(img_path)
        label = self.get_label_at(index)
        return img, label

    def __len__(self) -> int:
        return len(self.anno_df)
    
    def get_label_at(self, index: int) -> int:
        label = self.anno_df['level'].iloc[index].item()
        return label

    def get_path_at(self, index: int) -> str:
        img_name = self.anno_df['image'].iloc[index]
        img_path = os.path.join(self.data_path, "train", img_name+self.ext)
        return img_path


class Purpose(Enum):
    Train = 0
    Val = 1


FeatureAndTargetTransforms = Tuple[Callable[..., torch.Tensor],
                                   Callable[..., torch.Tensor]]

TensorRecord = Tuple[torch.Tensor, torch.Tensor]

def normalize(arr: np.ndarray) -> np.ndarray:
    return arr / np.sum(arr)


class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
    def __init__(self, dataset: RetinopathyDataset,
                 indices: np.ndarray,
                 purpose: Purpose,
                 transforms: FeatureAndTargetTransforms,
                 oversample_factor: int = 1,
                 stratify_classes: bool = False,
                 use_log_frequencies: bool = False,
                 ):

        self.dataset = dataset
        self.indices = indices
        self.purpose = purpose
        self.feature_transform = transforms[0]
        self.target_transform = transforms[1]
        self.oversample_factor = oversample_factor
        self.stratify_classes = stratify_classes
        self.use_log_frequencies = use_log_frequencies

        self.per_class_indices: Optional[Dict[int, np.ndarray]] = None
        self.frequencies: Optional[Dict[int, float]] = None
        if self.stratify_classes:
            self.bucketize_indices()
            if self.use_log_frequencies:
                self.calc_frequencies()

    def calc_frequencies(self):
        assert self.per_class_indices is not None
        counts_dict = {lbl: len(arr) for lbl, arr in self.per_class_indices.items()}
        counts = np.array(list(counts_dict.values()))
        counts_nrm = normalize(counts)
        temperature = 50.0 # > 1 to even-out frequencies
        freqs = normalize(np.log1p(counts_nrm * temperature))
        self.frequencies = {k: freq.item() for k, freq
                            in zip(self.per_class_indices.keys(), freqs)}
        print(self.frequencies)

    def bucketize_indices(self):
        buckets = defaultdict(list)
        for index in self.indices:
            label = self.dataset.get_label_at(index)
            buckets[label].append(index)
        self.per_class_indices = {k: np.array(v)
                                  for k, v in buckets.items()}

    def __getitem__(self, index: Union[int, slice]) -> TensorRecord: # type: ignore[override]
        assert isinstance(index, int)
        if self.purpose == Purpose.Train:
            index_rem = index % len(self.indices)
            idx = self.indices[index_rem].item()
        else:
            idx = self.indices[index].item()
        if self.per_class_indices:
            if self.frequencies is not None:
                arange = np.arange(len(self.per_class_indices))
                frequencies = np.zeros(len(self.per_class_indices), dtype=float)
                for k, v in self.frequencies.items():
                    frequencies[k] = v
                random_key = np.random.choice(
                    arange,
                    p=frequencies)
            else:
                random_key = np.random.randint(len(self.per_class_indices))

            indices = self.per_class_indices[random_key]
            actual_index = np.random.choice(indices).item()
        else:
            actual_index = idx
        feature, target = self.dataset[actual_index]
        feature_tensor = self.feature_transform(feature)
        target_tensor = self.target_transform(target)
        return feature_tensor, target_tensor

    def __len__(self):
        if self.purpose == Purpose.Train:
            return len(self.indices) * self.oversample_factor
        else:
            return len(self.indices)

    @staticmethod
    def make_splits(all_data: RetinopathyDataset,
                    train_transforms: FeatureAndTargetTransforms,
                    val_transforms: FeatureAndTargetTransforms,
                    train_fraction: float,
                    stratify_train: bool,
                    stratify_val: bool,
                    seed: int = 54,
                    ) -> Tuple['Split', 'Split']:

        prng = RandomState(seed)

        num_train = int(len(all_data) * train_fraction)
        all_indices = prng.permutation(len(all_data))
        train_indices = all_indices[:num_train]
        val_indices = all_indices[num_train:]
        train_data = Split(all_data, train_indices, Purpose.Train,
                           train_transforms, stratify_classes=stratify_train)
        val_data = Split(all_data, val_indices, Purpose.Val,
                         val_transforms, stratify_classes=stratify_val)
        return train_data, val_data


def print_data_stats(dataset: Union[Iterable[DataRecord], DataLoader], split_name: str) -> None:
    labels = []
    for _, label in dataset:
        if isinstance(label, torch.Tensor):
            label = label.cpu().numpy()
        labels.append(label)
    labels = np.concatenate(labels)
    cnt = Counter(labels)
    print(cnt)


class Metrics:
    def __init__(self,
                    num_classes: int,
                    labelmap: Dict[int, str],
                    split: str,
                    log_fn: Callable[..., None]) -> None:
        self.labelmap = labelmap
        self.loss = MeanMetric(nan_strategy='ignore')
        self.accuracy = MulticlassAccuracy(num_classes=num_classes)
        self.per_class_accuracies = MulticlassAccuracy(
            num_classes=num_classes, average=None)
        self.kappa = MulticlassCohenKappa(num_classes)
        self.split = split
        self.log_fn = log_fn
    
    def update(self,
               loss: torch.Tensor,
               preds: torch.Tensor,
               labels: torch.Tensor) -> None:
        self.loss.update(loss)
        self.accuracy.update(preds, labels)
        self.per_class_accuracies.update(preds, labels)
        self.kappa.update(preds, labels)

    def log(self) -> None:
        loss = self.loss.compute()
        accuracy = self.accuracy.compute()
        accuracies = self.per_class_accuracies.compute()
        kappa = self.kappa.compute()
        mean_accuracy = torch.nanmean(accuracies)
        self.log_fn(f"{self.split}/loss", loss, sync_dist=True)
        self.log_fn(f"{self.split}/accuracy", accuracy, sync_dist=True)
        self.log_fn(f"{self.split}/mean_accuracy", mean_accuracy, sync_dist=True)
        for i_class, acc in enumerate(accuracies):
            name = self.labelmap[i_class]
            self.log_fn(f"{self.split}/acc/{i_class} {name}", acc, sync_dist=True)
        self.log_fn(f"{self.split}/kappa", kappa, sync_dist=True)

    def to(self, device) -> 'Metrics':
        self.loss.to(device) # BUG HERE? should I assign it back?
        self.accuracy.to(device)
        self.per_class_accuracies.to(device)
        self.kappa.to(device)
        return self


def worker_init_fn(worker_id):
    state = np.random.get_state()
    assert isinstance(state, tuple)
    assert isinstance(state[1], np.ndarray)
    seed_arr = state[1]
    seed_np = seed_arr[0] + worker_id
    np.random.seed(seed_np)
    seed_pt = seed_np + 1111
    torch.manual_seed(seed_pt)
    print(f"Setting numpy seed to {seed_np} and pytorch seed to {seed_pt} in worker {worker_id}")


class ViTLightningModule(L.LightningModule):
    def __init__(self, debug: bool) -> None:
        super().__init__()

        self.save_hyperparameters()

        np.random.seed(53)

        # pretrained_name = 'google/vit-base-patch16-224-in21k'
        # pretrained_name = 'google/vit-base-patch16-384-in21k'

        # pretrained_name = "microsoft/resnet-50"
        pretrained_name = "microsoft/resnet-34"

        # processor = ViTImageProcessor.from_pretrained(pretrained_name)
        processor = AutoImageProcessor.from_pretrained(pretrained_name)

        image_mean = processor.image_mean # type: ignore
        image_std = processor.image_std # type: ignore
        # size = processor.size["height"] # type: ignore
        # size = processor.size["shortest_edge"] # type: ignore
        size = 896 # 448

        normalize = Normalize(mean=image_mean, std=image_std)
        train_transforms = Compose(
            [
                # RandomRotation((-180, 180)),
                RandomAffine((-180, 180), shear=10),
                RandomResizedCrop(size, scale=(0.5, 1.0)),
                RandomHorizontalFlip(),
                ToTensor(),
                normalize,
            ]
        )
        val_transforms = Compose(
            [
                Resize(size),
                CenterCrop(size),
                ToTensor(),
                normalize,
            ]
        )

        self.dataset = RetinopathyDataset("retinopathy_data")

        # print_data_stats(self.dataset, "all_data")

        train_data, val_data = Split.make_splits(
            self.dataset,
            train_transforms=(train_transforms, torch.tensor),
            val_transforms=(val_transforms, torch.tensor),
            train_fraction=0.9,
            stratify_train=True,
            stratify_val=True,
            )

        assert len(set(train_data.indices).intersection(set(val_data.indices))) == 0

        label2id = {label: id for id, label in self.dataset.label_map.items()}

        num_classes = len(self.dataset.label_map)
        labelmap = self.dataset.label_map
        assert len(labelmap) == num_classes
        assert set(labelmap.keys()) == set(range(num_classes))

        train_batch_size = 4 if debug else 20
        val_batch_size = 4 if debug else 20

        num_gpus = torch.cuda.device_count()
        print(f"{num_gpus=}")

        num_cores = torch.get_num_threads()
        print(f"{num_cores=}")

        num_threads_per_gpu = max(1, int(math.ceil(num_cores / num_gpus))) \
            if num_gpus > 0 else 1

        num_workers = 1 if debug else num_threads_per_gpu
        print(f"{num_workers=}")

        self._train_dataloader = DataLoader(
            train_data,
            shuffle=True,
            num_workers=num_workers,
            persistent_workers=num_workers > 0,
            pin_memory=True,
            batch_size=train_batch_size,
            worker_init_fn=worker_init_fn,
            )
        self._val_dataloader = DataLoader(
            val_data,
            shuffle=False,
            num_workers=num_workers,
            persistent_workers=num_workers > 0,
            pin_memory=True,
            batch_size=val_batch_size,
            )

        # print_data_stats(self._val_dataloader, "val")
        # print_data_stats(self._train_dataloader, "train")

        img_batch, label_batch = next(iter(self._train_dataloader))
        assert isinstance(img_batch, torch.Tensor)
        assert isinstance(label_batch, torch.Tensor)
        print(f"{img_batch.shape=} {label_batch.shape=}")
        
        assert img_batch.shape == (train_batch_size, 3, size, size)
        assert label_batch.shape == (train_batch_size,)
        
        self.example_input_array = torch.randn_like(img_batch)

        # self._model = ViTForImageClassification.from_pretrained(
        #     pretrained_name,
        #     num_labels=len(self.dataset.label_map),
        #     id2label=self.dataset.label_map,
        #     label2id=label2id)

        self._model = ResNetForImageClassification.from_pretrained(
            pretrained_name,
            num_labels=len(self.dataset.label_map),
            id2label=self.dataset.label_map,
            label2id=label2id,
            ignore_mismatched_sizes=True)

        assert isinstance(self._model, nn.Module)

        self.train_metrics: Optional[Metrics] = None
        self.val_metrics: Optional[Metrics] = None

    @property
    def num_classes(self):
        return len(self.dataset.label_map)
    
    @property
    def labelmap(self):
        return self.dataset.label_map

    def forward(self, img_batch):
        outputs = self._model(img_batch) # type: ignore
        return outputs.logits
        
    def common_step(self, batch, batch_idx):
        img_batch, label_batch = batch

        logits = self(img_batch)

        criterion = nn.CrossEntropyLoss()
        loss = criterion(logits, label_batch)
        preds_batch = logits.argmax(-1)

        return loss, preds_batch, label_batch

    def on_train_epoch_start(self) -> None:
        self.train_metrics = Metrics(
            self.num_classes,
            self.labelmap,
            "train",
            self.log).to(self.device)

    def training_step(self, batch, batch_idx):
        loss, preds, labels = self.common_step(batch, batch_idx)
        assert self.train_metrics is not None
        self.train_metrics.update(loss, preds, labels)

        if False and batch_idx == 0:
            self._dump_train_images()

        return loss

    def _dump_train_images(self) -> None:
        img_batch, label_batch = next(iter(self._train_dataloader))
        for i_img, (img, label) in enumerate(zip(img_batch, label_batch)):
            img_np = img.cpu().numpy()
            denorm_np = (img_np - img_np.min()) / (img_np.max() - img_np.min())
            img_uint8 = (255 * denorm_np).astype(np.uint8)
            pil_img = Image.fromarray(np.transpose(img_uint8, (1, 2, 0)))
            if self.logger is not None and self.logger.log_dir is not None:
                assert isinstance(self.logger.log_dir, str)
                os.makedirs(self.logger.log_dir, exist_ok=True)
                path = os.path.join(self.logger.log_dir,
                                    f"img_{i_img:02d}_{label.item()}.png")
                pil_img.save(path)

    def on_train_epoch_end(self) -> None:
        assert self.train_metrics is not None
        self.train_metrics.log()
        assert self.logger is not None
        if self.logger.log_dir is not None:
            path = os.path.join(self.logger.log_dir, "inference")
            self.save_checkpoint_dk(path)
    
    def save_checkpoint_dk(self, dirpath: str) -> None:
        if self.global_rank == 0:
            self._model.save_pretrained(dirpath)

    def validation_step(self, batch, batch_idx):
        loss, preds, labels = self.common_step(batch, batch_idx)
        assert self.val_metrics is not None
        self.val_metrics.update(loss, preds, labels)
        return loss

    def on_validation_epoch_start(self) -> None:
        self.val_metrics = Metrics(
            self.num_classes,
            self.labelmap,
            "val",
            self.log).to(self.device)
    
    def on_validation_epoch_end(self) -> None:
        assert self.val_metrics is not None
        self.val_metrics.log()

    def configure_optimizers(self):
        # No WD is the same as 1e-3 and better than 1e-2
        # LR 1e-3 is worse than 1e-4 (without LR scheduler)
        return AdamW(self.parameters(),
                     lr=1e-4,
                     )


def main():

    parser = ArgumentParser(description='KAUST-SDAIA Diabetic Retinopathy')
    parser.add_argument('--tag', action='store', type=str,
                        help='Extra suffix to put on the artefact dir name')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--convert-checkpoint', action='store', type=str,
                        help='Convert a checkpoint from training to pickle-independent '
                             'predictor-compatible directory')

    args = parser.parse_args()


    torch.set_float32_matmul_precision('high') # for V100/A100

    if args.convert_checkpoint is not None:

        print("Converting checkpoint", args.convert_checkpoint)

        checkpoint = torch.load(args.convert_checkpoint, map_location="cpu")
        print(list(checkpoint.keys()))

        model = ViTLightningModule.load_from_checkpoint(
            args.convert_checkpoint,
            map_location="cpu",
            hparams_file="tmp_ckpt_deleteme.yaml")

        model.save_checkpoint_dk("tmp_checkp_path_deleteme")

        print("Saved checkpoint. Done.")

    else:

        print("Start training")

        fast_dev_run = True if args.debug == True else False

        model = ViTLightningModule(fast_dev_run)

        datetime_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        art_dir_name = (f"{datetime_str}" +
                        (f"_{args.tag}" if args.tag is not None else ""))
        logger = TensorBoardLogger(save_dir=".", name="lightning_logs", version=art_dir_name)

        trainer = Trainer(
            logger=logger,
            benchmark=True,
            devices="auto",
            accelerator="auto",
            max_epochs=-1,
            callbacks=[
                ModelSummary(max_depth=-1),
                ],
            fast_dev_run=fast_dev_run,
            log_every_n_steps=10,
            )

        trainer.fit(
            model,
            train_dataloaders=model._train_dataloader,
            val_dataloaders=model._val_dataloader,
            )

        print("Training done")


if __name__ == "__main__":
    main()