finetune_cdmam_model_pl4.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL']= '3'
os.environ['CUDA_VISIBLE_DEVICES']= ''  # block TF CUDA init

import sys
import logging
import cv2
import glob
import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
from torch.optim import SGD
import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping
from tqdm import tqdm
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import KFold

#----- CUDA benchmarking optimization for fixed-size inputs

torch.backends.cudnn.benchmark= True  # let cuDNN to find optimal algorithms

#----- silence Pytorch Lightning

logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
os.environ['CUDA_VISIBLE_DEVICES']= '0'  # re-enable for Pytorch
torch.set_float32_matmul_precision('medium')

CLASSES= ['q0', 'q1', 'q2', 'q3']
NE= 300
BS= 32
LR= 1e-5
ES= 20

#----- command line arguments

img_path= sys.argv[1]
writ_app= sys.argv[2]
f_output= sys.argv[3]
mdl_name= sys.argv[4]

class ResNet18Custom(pl.LightningModule):
    
    def __init__(self, num_classes, lr= 1e-5, weight_decay= 1e-2, max_lr= 1e-03):
        
        super().__init__()
        self.save_hyperparameters()
        
        # base= models.resnet18(weights= models.ResNet18_Weights.DEFAULT)
        # base.fc= nn.Sequential(nn.Linear(base.fc.in_features, 128),
        #                        nn.ReLU(),
        #                        nn.Dropout(0.5),
        #                        nn.Linear(128, num_classes))
        
        base= models.resnet18(weights= models.ResNet18_Weights.DEFAULT)
        base.fc= nn.Sequential(nn.Linear(base.fc.in_features, 64),
                               nn.ReLU(),
                               nn.Dropout(0.5),
                               nn.Linear(64, num_classes))
        
        self.model= base
        
    def forward(self, x):
        
        return self.model(x)
    
class SimpleDataset(Dataset):
    
    def __init__(self, images, labels):
        
        self.images= torch.tensor(images.transpose(0, 3, 1, 2), dtype= torch.float32)
        self.labels= torch.tensor(labels, dtype= torch.float32)
        
    def __len__(self):
        
        return len(self.images)
    
    def __getitem__(self, idx):
        
        return self.images[idx], self.labels[idx]

class BlobDetector(pl.LightningModule):
    
    def __init__(self, num_classes= len(CLASSES), lr= 1e-5, weight_decay= 1e-2):
        
        super().__init__()
        self.save_hyperparameters()
        
        self.model= ResNet18Custom(num_classes, lr= 1e-5, weight_decay= 1e-2, max_lr= 1e-03)
        
        ckpt= torch.load(mdl_name, map_location= self.device)
        state_dict= ckpt['state_dict']
        new_state_dict= {k[len('model.'): ] if k.startswith('model.') else k: v for k, v in state_dict.items()}
        self.model.model.load_state_dict(new_state_dict)
        
        for param in self.model.parameters():
            param.requires_grad= False
        for param in self.model.model.layer4[1].parameters():
            param.requires_grad= True
        for param in self.model.model.fc.parameters():
            param.requires_grad= True
            
        frozen= [n for n, p in self.model.named_parameters() if not p.requires_grad]
        trainable= [n for n, p in self.model.named_parameters() if p.requires_grad]
        
        print(f"Frozen:    {len(frozen)} params: {frozen}")
        print(f"Trainable: {len(trainable)} params: {trainable}")
        
        self.model= torch.compile(self.model)
        self.criterion= nn.CrossEntropyLoss()
        
    def forward(self, x):
        
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        
        images, labels= batch
        outputs= self(images)
        loss= self.criterion(outputs, labels.argmax(dim= 1))
        acc= (outputs.argmax(dim= 1)== labels.argmax(dim= 1)).float().mean()
        
        self.log("train_loss", loss, on_step= False, on_epoch= True)
        self.log("train_acc", acc, on_step= False, on_epoch= True)
        
        return loss
    
    def validation_step(self, batch, batch_idx):
        
        images, labels= batch
        outputs= self(images)
        loss= self.criterion(outputs, labels.argmax(dim= 1))
        acc= (outputs.argmax(dim= 1)== labels.argmax(dim= 1)).float().mean()
        
        self.log("val_loss", loss, on_step= False, on_epoch= True)
        self.log("val_acc", acc, on_step= False, on_epoch= True)
        
    def predict_step(self, batch, batch_idx):
        
        images, labels= batch
        outputs= self(images)
        probs= torch.softmax(outputs, dim= 1)
        preds= outputs.argmax(dim= 1)
        
        return probs, preds, labels.argmax(dim= 1)
    
    def configure_optimizers(self):
        
        return SGD([p for p in self.model.parameters() if p.requires_grad],
                   lr= self.hparams.lr,
                   momentum= 0.9,
                   weight_decay= self.hparams.weight_decay)

class DBTDataModule(pl.LightningDataModule):
    
    def __init__(self, train_data, train_labels, val_data, val_labels, batch_size= 32):
        
        super().__init__()
        self.train_data= train_data
        self.train_labels= train_labels
        self.val_data= val_data
        self.val_labels= val_labels
        self.batch_size= batch_size
        self.num_workers= 4
        
    def setup(self, stage= None):
        
        self.train_dataset= SimpleDataset(self.train_data, self.train_labels)
        self.val_dataset= SimpleDataset(self.val_data, self.val_labels)
        
    def train_dataloader(self):
        
        return DataLoader(self.train_dataset,
                          batch_size= self.batch_size,
                          shuffle= True,
                          num_workers= self.num_workers,
                          pin_memory= True,
                          persistent_workers= True,
                          prefetch_factor= 2)

    def val_dataloader(self):
        
        return DataLoader(self.val_dataset,
                          batch_size= self.batch_size,
                          shuffle= False,
                          num_workers= self.num_workers,
                          pin_memory= True,
                          persistent_workers= True,
                          prefetch_factor= 2)

#----- wrap execution into main()

def main():
    
    print('[INFO] loading images...')
    
    data, labl, diam= [], [], []
    
    for f in tqdm(glob.glob(img_path+'/*.png')):
        
        lbl= f[-6:].split('.')[0]
        dia= f[-18:-7]
        img= cv2.imread(f, -1)
        im3= np.repeat(img[..., np.newaxis], 3, axis= -1)/img.max()
        im3= cv2.resize(im3, (224, 224))
        
        data.append(im3)
        labl.append(lbl)
        diam.append(dia)
        
    data= np.array(data, dtype= 'float32')
    labl= np.array(labl)
    diam= np.array(diam)
    
    lb= LabelBinarizer()
    labl_hot= lb.fit_transform(labl)
    labl= (diam, labl_hot)
    
    kf= KFold(n_splits= 10, shuffle= True, random_state= None)
    pc_arr= []
    
    for train_idx, test_idx in kf.split(data, labl[1]):
        
        trainX, testX= data[train_idx], data[test_idx]
        trainY, testY= labl[1][train_idx], labl[1][test_idx]
        testZ= labl[0][test_idx]
        
        dm= DBTDataModule(trainX, trainY, testX, testY, batch_size= BS)
        model= BlobDetector(num_classes= len(CLASSES), lr= LR)
        early_stopping= EarlyStopping(monitor= 'val_loss',
                                      patience= ES,
                                      mode= 'min',
                                      verbose= False)
        
        trainer= pl.Trainer(max_epochs= NE,
                            accelerator= 'gpu',
                            strategy= 'auto',
                            devices= 1,
                            callbacks= [early_stopping],
                            precision= '16-mixed',
                            logger= False,
                            enable_progress_bar= False,
                            enable_model_summary= False,
                            enable_checkpointing= False,
                            val_check_interval= 1.0,
                            default_root_dir= None,
                            gradient_clip_val= None,
                            profiler= None,
                            num_sanity_val_steps= 0)
        
        print('[INFO] training network...')
        trainer.fit(model, dm)
        
        print('[INFO] evaluating network...')
        predictions= trainer.predict(model, dm.val_dataloader())
        
        n_hits= 0
        for probs, preds, labels in predictions:
            n_hits+= (preds== labels).sum().item()
            
        pc= n_hits/len(testX)
        pc_arr.append(pc)
        
        print(f"4-AFC PC: {pc:.3f}")
        
    pc_avg= np.mean(pc_arr)
    pc_std= np.std(pc_arr)
    
    print(f"Mean PC: {pc_avg:.3f} +/- {pc_std:.3f}")
    
    f_res= open(f_output, writ_app)
    f_res.write(" ".join([f"{x:.3f}" for x in pc_arr])+'\n')
    f_res.close()
    
if __name__ == "__main__":
    
    main()