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]):
    """ A class to access the pre-downloaded Diabetic Retinopathy dataset. """

    def __init__(self, data_path: str) -> None:
        """ Constructor.

        Args:
            data_path (str): path to the dataset, ex: "retinopathy_data"
                containing "trainLabels.csv" and "train/".
        """
        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


""" Purpose of a split: training or validation. """
class Purpose(Enum):
    Train = 0
    Val = 1

""" Augmentation transformations for an image and a label. """
FeatureAndTargetTransforms = Tuple[Callable[..., torch.Tensor],
                                   Callable[..., torch.Tensor]]

""" Feature (image) and target (label) tensors. """
TensorRecord = Tuple[torch.Tensor, torch.Tensor]


class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]):
    """ Split is a class that keep a view on a part of a dataset.
    Split is used to hold the imormation about which samples go to training
    and which to validation without a need to put these groups of files into
    separate folders.
    """
    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,
                 ):
        """ Constructor.

        Args:
            dataset (RetinopathyDataset): The dataset on which the Split "views".
            indices (np.ndarray): Externally provided indices of samples that
                are "viewed" on.
            purpose (Purpose): Either train or val, to be able to replicate
                the data for train split for effecient workers utilization.
            transforms (FeatureAndTargetTransforms): Functors of feature and
                target transforms.
            oversample_factor (int, optional): Expand the training dataset by
                replication to avoid dataloader stalls on epoch ends. Defaults to 1.
            stratify_classes (bool, optional): Whether to apply stratified sampling.
                Defaults to False.
            use_log_frequencies (bool, optional): If stratify_classes=True,
                whether to use logarithmic sampling strategy. If False, apply
                regular even sampling. Defaults to 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 = self._normalize(counts)
        temperature = 50.0 # > 1 to even-out frequencies
        freqs = self._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)

    @staticmethod
    def _normalize(arr: np.ndarray) -> np.ndarray:
        return arr / np.sum(arr)

    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']:

        """ Prepare train and val splits deterministically.

        Returns:
            Tuple[Split, Split]:
                - Train split
                - Val 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: int) -> None:
    """ Initialize workers in a way that they draw different
    random samples and do not repeat identical pseudorandom
    sequences of each other, which may be the case with Fork
    multiprocessing.

    Args:
        worker_id (int): id of a preprocessing worker process launched
        by one DDP training process. 
    """
    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):
    """ Lightning Module that implements neural network training hooks. """
    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:
        """ Save augmented images to disk for inspection. """
        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():
    """ Neural network trainer entry point. """

    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',
                        help="Dummy training cycle for testing purposes")
    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()