DmitriiKhizbullin's picture
Initial implementation of trainer and gradio app
699342a
raw
history blame
19.5 kB
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()