classifier-train.py

import os
import hydra
import torch
import pytorch_lightning
import numpy as np
from sklearn.metrics import confusion_matrix, classification_report
from sklearn.utils.multiclass import unique_labels
from pytorch_lightning.callbacks import ModelCheckpoint
from src.models.classifier import OralClassifierModule
from src.datasets.classifier_datamodule import OralClassificationDataModule
from src.loss_log import LossLogCallback
from src.utils import *


@hydra.main(version_base=None, config_path="./config", config_name="config_classification")
def main(cfg):

    if cfg.train.seed == -1:
        random_data = os.urandom(4)
        seed = int.from_bytes(random_data, byteorder="big")
        cfg.train.seed = seed
    torch.manual_seed(cfg.train.seed)

    callbacks = list()
    callbacks.append(get_early_stopping(cfg))
    callbacks.append(LossLogCallback(cfg))
    loggers = get_loggers(cfg)

    torch.set_float32_matmul_precision("high")

    # model
    model = OralClassifierModule(
        model=cfg.model.name,
        weights=cfg.model.weights,
        num_classes=cfg.model.num_classes,
        lr=cfg.train.lr,
        max_epochs=cfg.train.max_epochs,
        features_size=cfg.model.features_size,
        frozen_layers=cfg.train.frozen_layers
    )

    # datasets and transformations
    train_img_tranform, val_img_tranform, test_img_tranform, img_tranform = get_transformations(cfg)
    data = OralClassificationDataModule(
        train=cfg.dataset.train,
        val=cfg.dataset.val,
        test=cfg.dataset.test,
        batch_size=cfg.train.batch_size,
        train_transform = train_img_tranform,
        val_transform = val_img_tranform,
        test_transform = test_img_tranform,
        transform = img_tranform,
    )

    checkpoint_callback = ModelCheckpoint(
        save_top_k=1, 
        monitor="val_loss",
        mode="min"
        )
    callbacks.append(checkpoint_callback)
    # training
    trainer = pytorch_lightning.Trainer(
        logger=loggers,
        callbacks=callbacks,
        accelerator=cfg.train.accelerator,
        devices=cfg.train.devices,
        log_every_n_steps=1,
        max_epochs=cfg.train.max_epochs,
        #gradient_clip_val=0.1, 
        #gradient_clip_algorithm="value"
    )
    trainer.fit(model, data)

    trainer.test(dataloaders=data.test_dataloader(), ckpt_path='best')

    # prediction
    predictions = trainer.predict(model, data)
    predictions = torch.cat(predictions, dim=0)
    predictions = torch.argmax(predictions, dim=1)
    gt = torch.cat([y for _, y in data.test_dataloader()], dim=0)

    print(classification_report(gt, predictions))

    class_names = np.array(['Neoplastic', 'Aphthous', 'Traumatic'])
    path = cfg.log.path + cfg.log.dir 
    log_dir = path + '/' + get_last_version(path)
    log_report(gt, predictions, classes=class_names, log_dir=log_dir)


if __name__ == "__main__":
    main()