train_with_accelerate_vision.py

import os
import pandas as pd
import numpy as np
import torch
from tqdm import tqdm
from transformers import ViTModel, ViTFeatureExtractor 
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs
from safetensors.torch import save_file, load_file
import torch
import random
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import argparse
from utils import custom_vit_vision, custom_vit_embed,create_train_val_dataloaders, image_path_refine, KorniaGPUAugmentation
import warnings
warnings.filterwarnings('ignore')

CFG = {
    'Image_size':518,
    'EPOCHS':100,
    'MIN_LR':5e-5,
    'MAX_LR':3e-4,
    'SEED':42,
    'Train_BS':192,
    'Valid_BS':192,
    'optimizer':'AdamW',
    'scheduler':"CosineAnnealingLR",
    'model_name':"vit-base-patch16-384"
}

def set_seed(seed_value):
    os.environ['PYTHONHASHSEED'] = str(seed_value)
    random.seed(seed_value)
    np.random.seed(seed_value)
    torch.manual_seed(seed_value)
    torch.cuda.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def train(seed):
    # If rad-dino is used
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
    
    # data load
    set_seed(seed)
    df = pd.read_csv('/data/mimic3_cxr_jpg/mimic-cxr-dataset.csv')
    df = df[(df.ViewPosition == 'PA') | (df.ViewPosition == 'AP')].drop_duplicates('study_id')
    df['ImagePath'] = df.apply(image_path_refine, axis=1)
    mimic_columns = ['Atelectasis','Cardiomegaly','Consolidation','Edema','Enlarged Cardiomediastinum','Fracture','Lung Lesion','Lung Opacity','No Finding','Pleural Effusion','Pleural Other','Pneumonia','Pneumothorax']

    # split
    df = df[df[mimic_columns].fillna(0).replace(-1, 0).sum(axis=1) != 0]
    train_df, temp_df = train_test_split(df, test_size=0.3, random_state=42, shuffle=True)
    val_df, test_df = train_test_split(temp_df, test_size=2/3, random_state=42, shuffle=True)
    train_df, val_df, test_df =  train_df.reset_index(drop=True), val_df.reset_index(drop=True), test_df.reset_index(drop=True)

    # DataLoader
    processor = ViTFeatureExtractor.from_pretrained('/data/models/vit-base-patch16-384')
    train_loader, valid_loader = create_train_val_dataloaders(train_df, val_df, label_type="classify", train_bs=CFG['Train_BS'], valid_bs=CFG['Valid_BS'], num_workers=20)

    # model load
    model = ViTModel.from_pretrained('/data/models/vit-base-patch16-384')
    model = custom_vit_vision(model)
    model = model.to(accelerator.device)

    # others load
    augment_tool = KorniaGPUAugmentation().to('cuda')
    optimizer = torch.optim.AdamW(model.parameters(), lr = CFG['MAX_LR'], betas=(0.9,0.999), eps=1e-6, weight_decay=0.01, amsgrad=False)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,eta_min=CFG['MIN_LR'], T_max=len(train_loader))

    model, optimizer, train_loader, valid_loader, scheduler = accelerator.prepare(
        model, optimizer, train_loader, valid_loader, scheduler
    )

    criterion = torch.nn.BCEWithLogitsLoss()
    best_loss = 99
    df = pd.DataFrame(columns=['epoch','train_loss','valid_loss'])
    for epoch in range(1, CFG["EPOCHS"] + 1):
        i = 0
        model.train()
        train_loss = []
        for imgs, labels in tqdm(iter(train_loader)):
            pixel_values = augment_tool(imgs, True)
            # labels = torch.tensor(labels).to(accelerator.device)
            optimizer.zero_grad()
            output = model(pixel_values)
            loss = criterion(output, labels)
            accelerator.backward(loss)
            optimizer.step()
            train_loss.append(loss.clone().detach().cpu().numpy())
            scheduler.step()
            # i += 1
            # if i == 10:
            #     break
        _train_loss = np.mean(train_loss)
        model.eval()
        valid_loss = []
        i = 0
        with torch.no_grad():
            for imgs, labels in tqdm(iter(valid_loader)):
                pixel_values = augment_tool(imgs, False)
                # labels = torch.tensor(labels).to(accelerator.device)
                output = model(pixel_values)
                loss = criterion(output, labels)
                valid_loss.append(loss.clone().detach().cpu().numpy())
                # i += 1
                # if i == 10:
                #     break
        _val_loss = np.mean(valid_loss)
        accelerator.print(
            f"Epoch [{epoch}], Train Loss : [{_train_loss:.5f}] Val Loss : [{_val_loss:.5f}]]"
        )
        accelerator.wait_for_everyone()
        df = pd.concat([pd.DataFrame([[epoch,_train_loss,_val_loss]], columns=df.columns),df],ignore_index=True)
        df.to_csv(f'/data/code/CXR_embedding_research/history/vit-history-vision-{seed}-padding-aug-tmp.csv',index=False)
        if best_loss >= _val_loss:
            best_loss = _val_loss
            accelerator.save_model(model, f"/data/mimic_ckp/vit-base-patch16-384-vision-{seed}-padding-aug-tmp.pt")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the main function with a specific seed.")
    parser.add_argument('--seed', type=int, required=True, help="Seed value for random number generation")
    args = parser.parse_args()
    train(args.seed)