Initial implementation of trainer and gradio app

.gitattributes

@@ -0,0 +1,3 @@
+media/* filter=lfs diff=lfs merge=lfs -text
+release_ckpts/** filter=lfs diff=lfs merge=lfs -text
+demo_data/** filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,160 +1,10 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
+.vscode/
 
-# C extensions
-*.so
+/retinopathy_data
+/slurm_logs
 
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
+/lightning_logs/
 
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
+__pycache__
 
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+/dcgm/

app.py

@@ -0,0 +1,123 @@
+import os
+import gradio as gr
+import numpy as np
+import torch
+from typing import Tuple, Optional, Dict, List
+import glob
+from collections import defaultdict
+
+from transformers import (AutoImageProcessor,
+                          ResNetForImageClassification)
+
+from labelmap import DR_LABELMAP
+
+
+class App:
+    def __init__(self) -> None:
+
+        ckpt_name = "2023-12-24_20-02-18_30345221_V100_x4_resnet34/"
+
+        path = f"release_ckpts/{ckpt_name}/inference/"
+
+        self.image_processor = AutoImageProcessor.from_pretrained(path)
+
+        self.model = ResNetForImageClassification.from_pretrained(path)
+
+        example_lists = self._load_example_lists()
+
+        device = 'GPU' if torch.cuda.is_available() else 'CPU'
+
+        css = ".output-image, .input-image, .image-preview {height: 600px !important}"
+
+        with gr.Blocks(css=css) as ui:
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Row():
+                        predict_btn = gr.Button("Predict", size="lg")
+                    with gr.Row():
+                        gr.Markdown(f"Running on {device}")
+                with gr.Column(scale=4):
+                    # output = gr.Textbox(label="Retinopathy level prediction")
+                    output = gr.Label(num_top_classes=len(DR_LABELMAP),
+                                      label="Retinopathy level prediction")
+                with gr.Column(scale=4):
+                    gr.Markdown("![](https://media.githubusercontent.com/media/Obs01ete/retinopathy/master/media/logo1.png)")
+            with gr.Row():
+                with gr.Column(scale=9, min_width=100):
+                    image = gr.Image(label="Retina scan")
+                with gr.Column(scale=1, min_width=150):
+                    for cls_id in range(len(example_lists)):
+                        label = DR_LABELMAP[cls_id]
+                        with gr.Tab(f"{cls_id} : {label}"):
+                            gr.Examples(
+                                example_lists[cls_id],
+                                inputs=[image],
+                                outputs=[output],
+                                fn=self.predict,
+                                examples_per_page=10,
+                                run_on_click=True)
+
+            predict_btn.click(
+                fn=self.predict,
+                inputs=image,
+                outputs=output,
+                api_name="predict")
+        
+        self.ui = ui
+
+    def launch(self) -> None:
+        self.ui.queue().launch(share=True)
+
+    def predict(self, image: Optional[np.ndarray]):
+        if image is None:
+            return dict()
+        cls_name, prob, probs = self._infer(image)
+        message = f"Predicted class={cls_name}, prob={prob:.3f}"
+        print(message)
+        probs_dict = {f"{i} - {DR_LABELMAP[i]}": float(v)
+                      for i, v in enumerate(probs)}
+        return probs_dict
+    
+    def _infer(self, image_chw: np.ndarray) -> Tuple[str, float, np.ndarray]:
+        assert isinstance(self.model, ResNetForImageClassification)
+
+        inputs = self.image_processor(image_chw, return_tensors="pt")
+
+        with torch.no_grad():
+            output = self.model(**inputs)
+
+        logits_batch = output.logits
+        assert len(logits_batch.shape) == 2
+        assert logits_batch.shape[0] == 1
+        logits = logits_batch[0]
+        probs = torch.softmax(logits, dim=-1)
+        predicted_label = int(probs.argmax(-1).item())
+        prob = probs[predicted_label].item()
+        cls_name = self.model.config.id2label[predicted_label]
+        return cls_name, prob, probs.numpy()
+
+    @staticmethod
+    def _load_example_lists() -> Dict[int, List[str]]:
+
+        example_flat_list = glob.glob("demo_data/train/**/*.jpeg")
+
+        example_lists: Dict[int, List[str]] = defaultdict(list)
+        for path in example_flat_list:
+            dir, _ = os.path.split(path)
+            _, subdir = os.path.split(dir)
+            try:
+                cls_id = int(subdir)
+            except ValueError:
+                print(f"Cannot parse path {path}")
+                continue
+            example_lists[cls_id].append(path)
+        return example_lists
+
+
+def main():
+    app = App()
+    app.launch()
+
+
+if __name__ == "__main__":
+    main()

labelmap.py

@@ -0,0 +1,7 @@
+DR_LABELMAP = {
+    0: 'No DR',
+    1: 'Mild',
+    2: 'Moderate',
+    3: 'Severe',
+    4: 'Proliferative DR',
+}

mypy.ini

@@ -0,0 +1,3 @@
+[mypy]
+ignore_missing_imports = True
+check_untyped_defs = True

tag.sh

@@ -0,0 +1,2 @@
+TAG=stratval_24h
+EXTRA_KEY=

train.py

@@ -0,0 +1,563 @@
+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()

train_a100_x1.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH -N 1
+#SBATCH -J SDAIA_DR
+#SBATCH -o slurm_logs/output.%J.out
+#SBATCH -e slurm_logs/output.%J.err
+#SBATCH --mail-user=${USER}@kaust.edu.sa
+#SBATCH --mail-type=END,FAIL
+#SBATCH --time=24:00:00
+#SBATCH --mem=100G
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:a100:1
+#SBATCH --cpus-per-task=16
+
+module purge
+source activate retinopathy
+
+echo Running four user "${USER}"
+
+. ./tag.sh
+
+PYTHONPATH=. python train.py --tag=\${SLURM_JOB_ID}_A100_x1_\${TAG}
+
+exit 0
+EOT

train_a100_x4.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH -N 1
+#SBATCH -J SDAIA_DR
+#SBATCH -o slurm_logs/output.%J.out
+#SBATCH -e slurm_logs/output.%J.err
+#SBATCH --mail-user=${USER}@kaust.edu.sa
+#SBATCH --mail-type=END,FAIL
+#SBATCH --time=24:00:00
+#SBATCH --mem=200G
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:a100:4
+#SBATCH --cpus-per-task=64
+
+module purge
+source activate retinopathy
+
+echo Running four user "${USER}"
+
+. ./tag.sh
+
+PYTHONPATH=. python train.py --tag=\${SLURM_JOB_ID}_A100_x4_\${TAG}
+
+exit 0
+EOT

train_v100_x1.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH -N 1
+#SBATCH -J SDAIA_DR
+#SBATCH -o slurm_logs/output.%J.out
+#SBATCH -e slurm_logs/output.%J.err
+#SBATCH --mail-user=${USER}@kaust.edu.sa
+#SBATCH --mail-type=END,FAIL
+#SBATCH --time=24:00:00
+#SBATCH --mem=100G
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:v100:1
+#SBATCH --cpus-per-task=10
+
+module purge
+source activate retinopathy
+
+echo Running four user "${USER}"
+
+. ./tag.sh
+
+PYTHONPATH=. python train.py --tag=\${SLURM_JOB_ID}_V100_x1_\${TAG}
+
+exit 0
+EOT

train_v100_x4.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+sbatch <<EOT
+#!/bin/bash
+
+#SBATCH -N 1
+#SBATCH -J SDAIA_DR
+#SBATCH -o slurm_logs/output.%J.out
+#SBATCH -e slurm_logs/output.%J.err
+#SBATCH --mail-user=${USER}@kaust.edu.sa
+#SBATCH --mail-type=END,FAIL
+#SBATCH --time=24:00:00
+#SBATCH --mem=200G
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:v100:4
+#SBATCH --cpus-per-task=40
+
+module purge
+source activate retinopathy
+
+echo Running four user "${USER}"
+
+. ./tag.sh
+
+PYTHONPATH=. python train.py --tag=\${SLURM_JOB_ID}_V100_x4_\${TAG}
+
+exit 0
+EOT