train.py

from DataGenerator import DatasetGenerator, ChexDataModule, CFG
from torchvision import transforms
from model import Densenet201
import pytorch_lightning as pl
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.callbacks import ModelCheckpoint


# ##  Data Directories 
# data_dir= r'C:\Users\patra\Desktop\New folder\ds' ##  Folder with images
# testfile= r'C:\Users\patra\Desktop\New folder\test.txt'
# trainfile= r'C:\Users\patra\Desktop\New folder\train.txt'
# valfile=r'C:\Users\patra\Desktop\New folder\val.txt'  

##Data Augmentation
train_transform = transforms.Compose([ 
    transforms.RandomRotation(20),
    transforms.RandomResizedCrop(224, scale=(0.63, 1)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

test_transform = transforms.Compose([ 
    # transforms.Resize(230),
    transforms.CenterCrop(224),
    transforms.ToTensor(), 
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

##Data Generators
train_data = DatasetGenerator(pathImageDirectory=CFG.data_dir, pathDatasetFile=CFG.trainfile, transform= train_transform)
test_data= DatasetGenerator(pathImageDirectory=CFG.data_dir, pathDatasetFile=CFG.testfile, transform=test_transform)
val_data= DatasetGenerator(pathImageDirectory= CFG.data_dir, pathDatasetFile=CFG.valfile, transform=test_transform)

def main():
    ## Lightning Data Module for Chexnet Images
    data_module= ChexDataModule(dataset=[train_data, val_data, test_data], batch_size=48, num_workers=16)
    logger= TensorBoardLogger("logs", name= 'chexnet', version= 1)
    checkpoint_callback= ModelCheckpoint(monitor='val_auroc', mode= 'max')
    trainer = pl.Trainer(accelerator="cuda", precision="16", logger=logger, callbacks=[checkpoint_callback], max_epochs=20)
    
    dense0= Densenet201()
    trainer.fit(dense0, data_module)


if __name__ == '__main__':
    main()