exp5/convnext2b_exp5_OBIDattention.py

from email.policy import strict
import os, time, pickle, shutil
import pandas as pd
import numpy as np

from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.cuda.amp import GradScaler
from torch import autocast

import torchvision.transforms as transforms

import timm
from timm.models import create_model
from timm.utils import ModelEmaV2

from timm.optim import create_optimizer_v2

from torchmetrics import MeanMetric
from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score
from torchmetrics import MetricCollection

from pytorch_metric_learning.losses import ArcFaceLoss

import wandb

import matplotlib.pyplot as plt


# ### parameters
################## Settings #############################
#os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
torch.backends.cudnn.benchmark = True

################## Data Paths ##########################
MODEL_DIR = "./convnext2b_obdid_attention/"

if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)
shutil.copyfile('./convnext2b_exp5_OBIDattention.py', f'{MODEL_DIR}convnext2b_exp5_OBIDattention.py')

TRAIN_DATA_DIR = "/SnakeCLEF2023-large_size/" # train imgs. path
ADD_TRAIN_DATA_DIR = "/HMP/" # add. train imgs. path 
VAL_DATA_DIR = "/SnakeCLEF2023-large_size/" # val imgs. path 

TRAINDATA_CONFIG = "/SnakeCLEF2023-TrainMetadata-iNat.csv"
ADD_TRAINDATA_CONFIG = "/SnakeCLEF2023-TrainMetadata-HM.csv"
VALIDDATA_CONFIG = "/SnakeCLEF2023-ValMetadata.csv"

MISSING_FILES = "../missing_train_data.csv" # csv with missing img. files that will be filtered out

CCM = "../code_class_mapping_obid.csv" # csv to metadata code to snake species dist.


NUM_CLASSES = 1784

################## Hyperparameters ########################
NUM_EPOCHS = 50
WARMUP_EPOCHS = 0
RESUME_EPOCH = 39 # resume model, optimizer from epoch 39 of experiment 4, checkpoint files need to be copied to the MODEL_DIR folder 


LEARNING_RATE = {
    'cnn': 1e-05,
    'embeddings': 1e-04,
    'classifier': 1e-04,
    'attention': 1e-04,
}

BATCH_SIZE = {
    'train': 1, 
    'valid': 1, 
    'grad_acc': 128, # gradient acc. steps with 'train' of batch sizes, global batch size = 'grad_acc' * 'train'
    'max_imgs_per_instance': 100 # maximum number of considered image instance (includes TTA) for each observation_id
}

BATCH_SIZE_AFTER_WARMUP = {
    'train': 1, 
    'valid': 1, 
    'grad_acc': 128, # gradient acc. steps with 'train' of batch sizes, global batch size = 'grad_acc' * 'train'
    'max_imgs_per_instance': 100 # maximum number of considered image instance (includes TTA) for each observation_id
}

TRANSFORMS = {
    'IMAGE_SIZE_TRAIN': 544,
    'IMAGE_SIZE_VAL': 544,
    'RandAug' : {
        'm': 7,
        'n': 2
    },
    'num_rand_crops': 5,  # num. of random crops during training per image instance
}


############# Focal Loss ####################
FOCAL_LOSS = {
    'class_dist': pickle.load(open("../classDist_HMP_missedRemoved.p", "rb"))['counts'],  # snake species frequency obtained on observation_id level taken into account missing observation_id of missing image files
    'gamma': 0.5,
}


############# Checkpoints ####################
CHECKPOINTS = {
    'fe_cnn': None,
    'model': None,
    'optimizer': None,
    'scaler': None,
    'arcloss': None,
    }

# ####### Embedding Token Mappings ########################
META_SIZES = {'endemic': 2, 'code': 212}
EMBEDDING_SIZES = {'endemic': 64, 'code': 64}

CODE_TOKENS = pickle.load(open("../meta_code_tokens.p", "rb"))
ENDEMIC_TOKENS = pickle.load(open("../meta_endemic_tokens.p", "rb"))

################### WandB ##################
WANDB = True

if WANDB:
    wandb.init(
        entity="snakeclef2023", # our team at wandb

        # set the wandb project where this run will be logged
        project="exp5", # -> define sub-projects here, e.g. experiments with MetaFormer or CNNs...

        # define a name for this run
        name="OBIDattention",

        # track all the used hyperparameters here, config is just a dict object so any key:value pairs are possible
        config={
            "learning_rate": LEARNING_RATE,
            "focal_loss": FOCAL_LOSS,
            "architecture": "convnextv2_base.fcmae_ft_in22k_in1k_384",
            "pretrained": "iNat21",
            "dataset": f"snakeclef2023, additional train data: {True if ADD_TRAINDATA_CONFIG else False}",
            "epochs": NUM_EPOCHS,
            "transforms": TRANSFORMS,
            "checkpoints": CHECKPOINTS,
            "model_dir": MODEL_DIR
            # ... any other hyperparameter that is necessary to reproduce the result
        },
        save_code=True, # save the script file as backup
        dir=MODEL_DIR # locally folder where wandb log files are saved
    )




##################### Dataset & AugTransforms #####################################
# ### dataset & loaders
class SnakeInstanceDataset(Dataset):
    def __init__(self, data, ccm, transform, fix_num=None):
        self.data = data
        self.instance_groups = data.groupby('observation_id').groups
        self.instance_obids = list(self.instance_groups.keys())
        
        self.transform = transform # Image augmentation pipeline
        self.code_class_mapping = ccm
        self.code_tokens = CODE_TOKENS
        self.endemic_tokens = ENDEMIC_TOKENS
        
        self.fix_num = fix_num
        self.random_gen = torch.Generator().manual_seed(1)


    def __len__(self):
        return len(self.instance_obids)

    def __getitem__(self, index):
        obid = self.instance_obids[index] # get observation id
        instances = self.data.iloc[self.instance_groups[obid]]
        
        code = instances.code.tolist()[0]
        code = code if code in self.code_tokens.keys() else "unknown"
        endemic = instances.endemic.tolist()[0] 
        endemic = endemic if endemic in self.endemic_tokens.keys() else False # get endemic metadata
        
        label = torch.tensor([instances.class_id.tolist()[0]]) # get "global" label 
        ccm = torch.from_numpy(self.code_class_mapping[code].to_numpy()) # code class mapping
        meta = torch.tensor([[self.code_tokens[code], self.endemic_tokens[endemic]]]) # metadata tokens
        
        # load instance images
        files = instances.image_path.tolist()
        imgs = torch.stack([self.transform(Image.open(file).convert("RGB")) for file in files])
        img_size = imgs.size(-1)
        imgs = imgs.view(-1, 3, img_size, img_size)
        
        # randomly shuffle imgs and/or draw subset of imgs
        num_imgs = imgs.size(0)
        idx = torch.randperm(num_imgs, generator=self.random_gen)
        idx = idx[:self.fix_num] if self.fix_num else idx # randomly draw 5 imgs 
        imgs = imgs[idx, :, :, :]
        
        return (imgs, label, ccm, meta)


# valid data preprocessing pipeline
def get_val_preprocessing(img_size):
    print(f'IMG_SIZE_VAL: {img_size}')
    return transforms.Compose([
                transforms.Resize(int(img_size * 1.25)), # Expand IMAGE_SIZE before random crop
                transforms.Compose([
                    transforms.FiveCrop((img_size, img_size)), # this is a list of PIL Images
                    transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])) # returns a 4D tensor
                ]),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ])

class MultipleRandomCropsWithAugmentation:
    def __init__(self, img_size, num_crops=5):
        super(MultipleRandomCropsWithAugmentation, self).__init__()
        self.num_crops = num_crops
        self.random_crop = transforms.RandomCrop((img_size, img_size))
        self.augment = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.RandomVerticalFlip(p=0.5),
                transforms.RandAugment(num_ops=TRANSFORMS['RandAug']['n'], magnitude=TRANSFORMS['RandAug']['m'])
        ])
        self.to_tensor = transforms.ToTensor()
        
    def __call__(self, x):
        x = torch.stack([self.to_tensor(self.augment(self.random_crop(x))) for i in range(self.num_crops)])           
        return x

# train data augmentation/ preprocessing pipeline
def get_train_augmentation_preprocessing(img_size):
    print(f'IMG_SIZE_TRAIN: {img_size}')
    return transforms.Compose([
                transforms.Resize(int(img_size * 1.25)), # Expand IMAGE_SIZE before random crop
                MultipleRandomCropsWithAugmentation(img_size, TRANSFORMS['num_rand_crops']),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ])


def get_datasets(train_transfroms, val_transforms):
    # load CSVs
    nan_values = ['', '#N/A', '#N/A N/A', '#NA', '-1.#IND', '-1.#QNAN', '-NaN', '-nan', '1.#IND', '1.#QNAN', '<NA>', 'N/A', 'NULL', 'NaN', 'n/a', 'nan', 'null']
    train_data = pd.read_csv(TRAINDATA_CONFIG, na_values=nan_values, keep_default_na=False)
    train_data = train_data.drop_duplicates(subset='image_path', keep="first")

    missing_train_data = pd.read_csv(MISSING_FILES, na_values=nan_values, keep_default_na=False)
    valid_data = pd.read_csv(VALIDDATA_CONFIG, na_values=nan_values, keep_default_na=False)
    valid_data = valid_data.drop_duplicates(subset='image_path', keep="first")


    # delete missing files of train data table
    train_data = pd.merge(train_data, missing_train_data, how='outer', indicator=True)
    train_data = train_data.loc[train_data._merge == 'left_only', ["observation_id","endemic","binomial_name","code","image_path","class_id","subset"]]

    # load transposed version of CCM table
    ccm = pd.read_csv(CCM, na_values=nan_values, keep_default_na=False)

    # add image path
    train_data["image_path"] = TRAIN_DATA_DIR + train_data['image_path']
    valid_data["image_path"] = VAL_DATA_DIR + valid_data['image_path']

    # add additional data
    if ADD_TRAINDATA_CONFIG:
        add_train_data = pd.read_csv(ADD_TRAINDATA_CONFIG, na_values=nan_values, keep_default_na=False)
        add_train_data["image_path"] = ADD_TRAIN_DATA_DIR + add_train_data['image_path']
        train_data = pd.concat([train_data, add_train_data], axis=0)
   
    # limit data size
    #train_data = train_data.head(150)
    #valid_data = valid_data.head(150)

    # shuffle
    train_data = train_data.sample(frac=1, random_state=1).reset_index(drop=True)
    valid_data = valid_data.sample(frac=1, random_state=1).reset_index(drop=True)

    # compute train, valid data weights
    #TCLASS_WEIGHTS = compute_weights(train_data)
    #VCLASS_WEIGHTS = compute_weights(valid_data)

    # create datasets
    train_dataset = SnakeInstanceDataset(train_data, ccm, transform=train_transfroms, fix_num=BATCH_SIZE['max_imgs_per_instance'])
    valid_dataset = SnakeInstanceDataset(valid_data, ccm, transform=val_transforms, fix_num=BATCH_SIZE['max_imgs_per_instance'])
    print(f'train dataset shape: {len(train_dataset)}')
    print(f'valid dataset shape: {len(valid_dataset)}')

    return train_dataset, valid_dataset#, TCLASS_WEIGHTS, VCLASS_WEIGHTS

def get_collate_fn():
    def collate_fn(batch):
        imgs = batch[0][0]
        targets = batch[0][1]
        ccm = batch[0][2]
        meta = batch[0][3]
        return [imgs, targets, ccm, meta]
    return collate_fn

def get_dataloaders(imgsize_train, imgsize_val):
    # get train, valid augmentation & preprocessing pipelines
    train_aug_preprocessing = get_train_augmentation_preprocessing(imgsize_train)
    val_preprocessing = get_val_preprocessing(imgsize_val)
    # prepare the datasets
    train_dataset, valid_dataset = get_datasets(train_transfroms=train_aug_preprocessing, val_transforms=val_preprocessing)
    train_loader = DataLoader(dataset=train_dataset, shuffle=True, batch_size=1, num_workers=4, prefetch_factor=8, collate_fn=get_collate_fn(), drop_last=False, pin_memory=True)
    valid_loader = DataLoader(dataset=valid_dataset, shuffle=False, batch_size=1, num_workers=4, prefetch_factor=8, collate_fn=get_collate_fn(), drop_last=False, pin_memory=True)

    return train_loader, valid_loader


# #################### plot train history #########################

def plot_history(logs):
    fig, ax = plt.subplots(3, 1, figsize=(8, 12))

    ax[0].plot(logs['loss'], label="train data")
    ax[0].plot(logs['val_loss'], label="valid data")
    ax[0].legend(loc="best")
    ax[0].set_ylabel("loss")
    ax[0].set_ylim([0, -np.log(1/NUM_CLASSES)])
    #ax[0].set_xlabel("epochs")
    ax[0].set_title("train- vs. valid loss")

    ax[1].plot(logs['acc'], label="train data")
    ax[1].plot(logs['val_acc'], label="valid data")
    ax[1].legend(loc="best")
    ax[1].set_ylabel("accuracy")
    ax[1].set_ylim([0, 1.01])
    #ax[1].set_xlabel("epochs")
    ax[1].set_title("train- vs. valid accuracy")

    ax[2].plot(logs['f1'], label="train data")
    ax[2].plot(logs['val_f1'], label="valid data")
    ax[2].legend(loc="best")
    ax[2].set_ylabel("f1")
    ax[2].set_ylim([0, 1.01])
    ax[2].set_xlabel("epochs")
    ax[2].set_title("train- vs. valid f1")

    fig.savefig(f'{MODEL_DIR}model_history.svg', dpi=150, format="svg")
    plt.show()

#################### Focal Loss ##################################
class FocalLoss(nn.Module):
    '''
    Multi-class Focal Loss
    '''
    def __init__(self, gamma=2, class_dist=None, reduction='mean', device='cuda'):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.weight = torch.tensor((1.0 - 0.999) / (1.0 - 0.999**class_dist), dtype=torch.float32, device=device) if class_dist is not None else torch.ones(NUM_CLASSES, device=device)
        self.reduction = reduction

    def forward(self, inputs, targets):
        """
        input: [N, C], float32
        target: [N, ], int64
        """
        logpt = torch.nn.functional.log_softmax(inputs, dim=1)
        pt = torch.exp(logpt)
        logpt = (1-pt)**self.gamma * logpt
        loss = torch.nn.functional.nll_loss(logpt, targets, weight=self.weight, reduction=self.reduction)
        return loss


# #################### Model #####################################

class FeatureExtractor(nn.Module):
    def __init__(self):
        super(FeatureExtractor, self).__init__()
        self.conv_backbone = create_model('convnextv2_base.fcmae_ft_in22k_in1k_384', pretrained=True, num_classes=0, drop_path_rate=0.2)
        if CHECKPOINTS['fe_cnn']:
            self.conv_backbone.load_state_dict(torch.load(CHECKPOINTS['fe_cnn'], map_location='cpu'), strict=True)
            print(f"use FE_CHECKPOINTS: {CHECKPOINTS['fe_cnn']}")
            torch.cuda.empty_cache()

    def forward(self, img):
        conv_features = self.conv_backbone(img)
        return conv_features


class MetaEmbeddings(nn.Module):
    def __init__(self, embedding_sizes: dict, meta_sizes: dict, dropout: float = None):
        super(MetaEmbeddings, self).__init__()
        self.endemic_embedding = nn.Embedding(meta_sizes['endemic'], embedding_sizes['endemic'], max_norm=1.0)
        self.code_embedding = nn.Embedding(meta_sizes['code'], embedding_sizes['code'], max_norm=1.0)

        self.dim_embedding = sum(embedding_sizes.values())
        self.embedding_net = nn.Sequential(
            nn.Linear(in_features=self.dim_embedding, out_features=self.dim_embedding, bias=True),
            nn.GELU(),
            nn.LayerNorm(self.dim_embedding, eps=1e-06),
            nn.Dropout(p=dropout, inplace=False) if dropout else nn.Identity(),
            nn.Linear(in_features=self.dim_embedding, out_features=self.dim_embedding, bias=True),
            nn.GELU(),
            nn.LayerNorm(self.dim_embedding, eps=1e-06),
        )

    def forward(self, meta):
        code_feature = self.code_embedding(meta[:,0])
        endemic_feature = self.endemic_embedding(meta[:,1])

        embeddings = torch.concat([code_feature, endemic_feature], dim=-1)
        embedding_features = self.embedding_net(embeddings)

        return embedding_features


class Classifier(nn.Module):
    def __init__(self, num_classes: int, dim_embeddings: int, dropout: float = None):
        super(Classifier, self).__init__()
        self.dropout = nn.Dropout(p=dropout, inplace=False) if dropout else nn.Identity()
        self.classifier = nn.Linear(in_features=dim_embeddings, out_features=num_classes, bias=True)

    def forward(self, embeddings):
        dropped_feature = self.dropout(embeddings)
        outputs = self.classifier(dropped_feature)

        return outputs

class Attention(nn.Module):
    def __init__(self):
        super(Attention, self).__init__()
        self.L = 1024
        self.D = 256
        self.K = 1

        self.attention = nn.Sequential(
            nn.Linear(self.L, self.D),
            nn.Tanh(),
            nn.Linear(self.D, self.K)
        )

    def forward(self, x):
        N, L = x.shape
        x = x.view(1,N,L)
                
        A = self.attention(x)  # 1xNx1
        A = torch.transpose(A, 2, 1)  # 1x1xN
        A = nn.functional.softmax(A, dim=-1)  # softmax over N
        M = torch.bmm(A, x).squeeze(dim=1)  # 1xL

        return M, A


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.feature_extractor = FeatureExtractor()
        self.embedding_net = MetaEmbeddings(embedding_sizes=EMBEDDING_SIZES, meta_sizes=META_SIZES, dropout=0.25)
        self.mil_pooling = Attention()
        self.classifier = Classifier(num_classes=NUM_CLASSES, dim_embeddings=1024+128, dropout=0.25)

    def forward(self, img, meta):
        img_features = self.feature_extractor(img)
        img_features, A = self.mil_pooling(img_features)

        meta_features = self.embedding_net(meta)
        cat_features = torch.concat([img_features, meta_features], dim=-1)
        classifier_outputs = self.classifier(cat_features)
        
        return classifier_outputs, cat_features

class LossLayer(nn.Module):
    def __init__(self):
        super(LossLayer, self).__init__()
        self.arcloss = ArcFaceLoss(num_classes=NUM_CLASSES, embedding_size=1024+128, margin=28.6, scale=64)
        self.celoss = FocalLoss(gamma=FOCAL_LOSS['gamma'], class_dist=FOCAL_LOSS['class_dist'])
        
    def forward(self, classifier_outputs, cat_features, labels):
        classifier_loss = self.celoss(classifier_outputs, labels)
        embedding_loss = self.arcloss(cat_features, labels)
        return classifier_loss + embedding_loss


def load_checkpoints(model=None, ema_model=None, optimizer=None, scaler=None, arcloss=None):
    if CHECKPOINTS['model'] and model is not None:
        model.load_state_dict(torch.load(CHECKPOINTS['model'], map_location='cpu'), strict=False)
        print(f"use model checkpoints: {CHECKPOINTS['model']}")
    if CHECKPOINTS['ema_model'] and ema_model is not None:
        ema_model.load_state_dict(torch.load(CHECKPOINTS['ema_model'], map_location='cpu'), strict=False)
        print(f"use ema_model checkpoints: {CHECKPOINTS['ema_model']}")
    if CHECKPOINTS['optimizer'] and optimizer is not None:
        optimizer.load_state_dict(torch.load(CHECKPOINTS['optimizer'], map_location='cpu'))
        print(f"use optimizer checkpoints: {CHECKPOINTS['optimizer']}")
    if CHECKPOINTS['scaler'] and scaler is not None:
        scaler.load_state_dict(torch.load(CHECKPOINTS['scaler'], map_location='cpu'))
        print(f"use scaler checkpoints: {CHECKPOINTS['scaler']}")
    if CHECKPOINTS['arcloss'] and arcloss is not None:
        arcloss.load_state_dict(torch.load(CHECKPOINTS['arcloss'], map_location='cpu'))
        print(f"use arcloss checkpoints: {CHECKPOINTS['arcloss']}")
    torch.cuda.empty_cache()

def resume_checkpoints(model=None, optimizer=None, scaler=None):
    if model is not None:
        model.load_state_dict(torch.load(f'{MODEL_DIR}model_epoch{RESUME_EPOCH}.pth', map_location='cpu'), strict=False)
        print(f"use model checkpoints: {MODEL_DIR}model_epoch{RESUME_EPOCH}.pth")
    if optimizer is not None:
        optimizer.load_state_dict(torch.load(f'{MODEL_DIR}optimizer_epoch{RESUME_EPOCH}.pth', map_location='cpu'))
        print(f"use optimizer checkpoints: {MODEL_DIR}optimizer_epoch{RESUME_EPOCH}.pth")
        
    if scaler is not None:
        scaler.load_state_dict(torch.load(f'{MODEL_DIR}mp_scaler_epoch{RESUME_EPOCH}.pth', map_location='cpu'))
        print(f"use scaler checkpoints: {MODEL_DIR}mp_scaler_epoch{RESUME_EPOCH}.pth")
    torch.cuda.empty_cache()


def resume_logs(logs):
    old_logs = pd.read_csv(f"{MODEL_DIR}train_history.csv")
    for m in list(logs.keys()):
        logs[m].extend(list(old_logs[m].values))

######################## Optimizer #####################################
def get_optm_group(module):
    """
    This long function is unfortunately doing something very simple and is being very defensive:
    We are separating out all parameters of the model into two buckets: those that will experience
    weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
    We are then returning the PyTorch optimizer object.
    """

    # separate out all parameters to those that will and won't experience regularizing weight decay
    decay = set()
    no_decay = set()
    whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv1d, timm.layers.GlobalResponseNormMlp)
    blacklist_weight_modules = (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.LayerNorm, torch.nn.Embedding)
    for mn, m in module.named_modules():
        for pn, p in m.named_parameters():
            fpn = '%s.%s' % (mn, pn) if mn else pn # full param name

            if pn.endswith('bias'):
                # all biases will not be decayed
                no_decay.add(fpn)
            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
                # weights of whitelist modules will be weight decayed
                decay.add(fpn)
            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
                # weights of blacklist modules will NOT be weight decayed
                no_decay.add(fpn)


    # validate that we considered every parameter
    param_dict = {pn: p for pn, p in module.named_parameters()}
    inter_params = decay & no_decay
    union_params = decay | no_decay
    assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
    assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
                                                % (str(param_dict.keys() - union_params), )

    return param_dict, decay, no_decay


def get_warmup_optimizer(model):
    params_group = []
    
    param_dict, decay, no_decay = get_optm_group(model.embedding_net)
    params_group.append({"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.05, 'lr': LEARNING_RATE['embeddings']})
    params_group.append({"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0, 'lr': LEARNING_RATE['embeddings']})

    param_dict, decay, no_decay = get_optm_group(model.classifier)
    params_group.append({"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.05, 'lr': LEARNING_RATE['classifier']})
    params_group.append({"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0, 'lr': LEARNING_RATE['classifier']})
    
    optimizer = torch.optim.AdamW(params_group)
    return optimizer


def get_after_warmup_optimizer(model, old_opt):
    new_opt = create_optimizer_v2(model.feature_extractor.conv_backbone, opt='adamw', filter_bias_and_bn=True, weight_decay=1e-8, layer_decay=0.85, lr=LEARNING_RATE['cnn'])
    
    # add old param groups
    for group in old_opt.param_groups:
        new_opt.add_param_group(group)
        
    return new_opt


# #################### Model Warmup #####################################

def warmup_start(model):
    # freeze model feature_extractor.conv_backbone during warmup
    for i, (param_name, param) in enumerate(model.feature_extractor.conv_backbone.named_parameters()):
        param.requires_grad = False
    print(f'--> freeze feature_extractor.conv_backbone during warmup phase')
    
    # freeze model feature_extractor.conv_backbone during warmup
    for i, (param_name, param) in enumerate(model.embedding_net.named_parameters()):
        param.requires_grad = False
    print(f'--> freeze feature_extractor.embedding_net during warmup phase')

def warmup_end(model):
    # unfreeze feature_extractor.conv_backbone during warmup
    for i, (param_name, param) in enumerate(model.feature_extractor.conv_backbone.named_parameters()):
        param.requires_grad = True
    print(f'--> unfreeze feature_extractor.conv_backbone after warmup phase')
    
    # freeze model feature_extractor.conv_backbone during warmup
    for i, (param_name, param) in enumerate(model.embedding_net.named_parameters()):
        param.requires_grad = True
    print(f'--> unfreeze feature_extractor.embedding_net during warmup phase')


# #################### Train Loop #####################################

# ### train
def main():
    device = torch.device(f'cuda:1')
    torch.cuda.set_device(device)

    # prepare the datasets
    train_loader, valid_loader = get_dataloaders(imgsize_train=TRANSFORMS['IMAGE_SIZE_TRAIN'],
                                                 imgsize_val=TRANSFORMS['IMAGE_SIZE_VAL'],
                                                 )

    # instantiate the model
    model = Model().to(device)
    if RESUME_EPOCH > 0:
        resume_checkpoints(model=model)
    ema_model = ModelEmaV2(model, decay=0.9998, device=device)
    warmup_start(model)
    
    loss_fn = LossLayer().to(device)
    if RESUME_EPOCH > 0:
        resume_checkpoints(arcloss=loss_fn.arcloss)

    # Optimizer & Schedules & early stopping
    optimizer = get_warmup_optimizer(model)
    optimizer.add_param_group({"params": loss_fn.arcloss.parameters(), "weight_decay": 0.0, 'lr': LEARNING_RATE['classifier']})

    scaler = GradScaler()
    if RESUME_EPOCH > 0:
        #optimizer = get_after_warmup_optimizer(model, optimizer) if RESUME_EPOCH > WARMUP_EPOCHS else optimizer
        resume_checkpoints(optimizer=optimizer, scaler=scaler)
    
    # add attention module
    param_dict, decay, no_decay = get_optm_group(model.attention)
    optimizer.add_param_group({"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.05, 'lr': LEARNING_RATE['attention']})
    optimizer.add_param_group({"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0, 'lr': LEARNING_RATE['attention']})

    # running metrics during training
    loss_metric = MeanMetric().to(device)
    metrics = MetricCollection(metrics={
        'acc': MulticlassAccuracy(num_classes=NUM_CLASSES, average='macro'),
        'top3_acc': MulticlassAccuracy(num_classes=NUM_CLASSES, average='macro', top_k=3),
        'f1': MulticlassF1Score(num_classes=NUM_CLASSES, average='macro')  
    }).to(device)
    metric_ccm = MulticlassF1Score(num_classes=NUM_CLASSES, average='macro').to(device)     

    # start time of trainig
    start_training = time.perf_counter()
    # create log dict
    logs = {'loss': [], 'acc': [], 'acc_top3': [], 'f1': [], 'f1country': [], 'val_loss': [], 'val_acc': [], 'val_acc_top3': [], 'val_f1': [], 'val_f1country': []}
    if RESUME_EPOCH > 0:
        resume_logs(logs)

    #iterate over epochs
    start_epoch = RESUME_EPOCH+1 if RESUME_EPOCH > 0 else 0
    for epoch in range(start_epoch, NUM_EPOCHS):
        # start time of epoch
        epoch_start = time.perf_counter()
        print(f'Epoch {epoch+1}/{NUM_EPOCHS}')
        
        ############################## train phase ####################################
        model.train()

        # zero the parameter gradients
        optimizer.zero_grad(set_to_none=True)
        
        # grad acc loss divider
        loss_div = torch.tensor(BATCH_SIZE['grad_acc'], dtype=torch.float16, device=device, requires_grad=False) if BATCH_SIZE['grad_acc'] != 0 else torch.tensor(1.0, dtype=torch.float16, device=device, requires_grad=False)
        
        # iterate over training batches
        for batch_idx, (inputs, labels, ccm, meta) in enumerate(train_loader):
            inputs = inputs.to(device, non_blocking=True)
            meta = meta.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            ccm = ccm.to(device, non_blocking=True) 
            
            # forward with mixed precision
            with autocast(device_type='cuda', dtype=torch.float16):
                outputs, embeddings = model(inputs, meta) 
                loss = loss_fn(outputs, embeddings, labels) / loss_div

            # loss backward 
            scaler.scale(loss).backward()
            
            # Compute metrics
            loss_metric.update((loss * loss_div).detach())
            
            preds = outputs.softmax(dim=-1).detach()
            metrics.update(preds, labels)
            metric_ccm.update(preds * ccm, labels)
            
            ############################ grad acc ##############################
            if (batch_idx+1) % BATCH_SIZE['grad_acc'] == 0:
                #scaler.unscale_(optimizer)
                #torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # optimize with gradient clipping to 1 with mixed precision
                scaler.step(optimizer)
                scaler.update()
                # zero the parameter gradients
                optimizer.zero_grad(set_to_none=True)
                # update ema model
                ema_model.update(model)
  

        # compute, sync & reset metrics for validation
        epoch_loss = loss_metric.compute()
        epoch_metrics = metrics.compute()
        epoch_metric_ccm = metric_ccm.compute()
        
        loss_metric.reset()
        metrics.reset()
        metric_ccm.reset()
        
        # Append metric results to logs
        logs['loss'].append(epoch_loss.cpu().item())
        logs['acc'].append(epoch_metrics['acc'].cpu().item())
        logs['acc_top3'].append(epoch_metrics['top3_acc'].cpu().item())
        logs['f1'].append(epoch_metrics['f1'].cpu().item())
        logs['f1country'].append(epoch_metric_ccm.detach().cpu().item())
        
        print(f"loss: {logs['loss'][epoch]:.5f}, acc: {logs['acc'][epoch]:.5f}, acc_top3: {logs['acc_top3'][epoch]:.5f}, f1: {logs['f1'][epoch]:.5f}, f1country: {logs['f1country'][epoch]:.5f}", end=' || ')
    
        # zero the parameter gradients
        optimizer.zero_grad(set_to_none=True)
        
        del inputs, labels, ccm, meta, preds, outputs, loss, loss_div, epoch_loss, epoch_metrics, epoch_metric_ccm
        torch.cuda.empty_cache()

        ############################## valid phase ####################################
        with torch.no_grad():
            model.eval()

            # iterate over validation batches
            for (inputs, labels, ccm, meta) in valid_loader:
                inputs = inputs.to(device, non_blocking=True)
                meta = meta.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                ccm = ccm.to(device, non_blocking=True)
                
                # forward with mixed precision
                with autocast(device_type='cuda', dtype=torch.float16):
                    outputs, embeddings = model(inputs, meta) 
                    loss = loss_fn(outputs, embeddings, labels)

                # Compute metrics
                loss_metric.update(loss.detach())

                preds = outputs.softmax(dim=-1).detach()
                metrics.update(preds, labels)
                metric_ccm.update(preds * ccm, labels)

            # compute, sync & reset metrics for validation
            epoch_loss = loss_metric.compute()
            epoch_metrics = metrics.compute()
            epoch_metric_ccm = metric_ccm.compute()

            loss_metric.reset()
            metrics.reset()
            metric_ccm.reset()

            # Append metric results to logs
            logs['val_loss'].append(epoch_loss.cpu().item())
            logs['val_acc'].append(epoch_metrics['acc'].cpu().item())
            logs['val_acc_top3'].append(epoch_metrics['top3_acc'].cpu().item())
            logs['val_f1'].append(epoch_metrics['f1'].cpu().item())
            logs['val_f1country'].append(epoch_metric_ccm.detach().cpu().item())
                                  
            print(f"val_loss: {logs['val_loss'][epoch]:.5f}, val_acc: {logs['val_acc'][epoch]:.5f}, val_acc_top3: {logs['val_acc_top3'][epoch]:.5f}, val_f1: {logs['val_f1'][epoch]:.5f}, val_f1country: {logs['val_f1country'][epoch]:.5f}", end=' || ')
            
        del inputs, labels, ccm, meta, preds, outputs, loss, epoch_loss, epoch_metrics, epoch_metric_ccm
        torch.cuda.empty_cache()
                                  
        # save logs as csv
        logs_df = pd.DataFrame(logs)
        logs_df.to_csv(f'{MODEL_DIR}train_history.csv', index_label='epoch', sep=',', encoding='utf-8')

        if WANDB:
            # at the end of each epoch, log anything you want to log for that epoch
            wandb.log(
                {k:v[epoch] for k,v in logs.items()}, # e.g. log each metric value for the current epoch in our defined logs dict
                step=epoch # epoch index for wandb
            )

        #save trained model for each epoch
        torch.save(model.state_dict(), f'{MODEL_DIR}model_epoch{epoch}.pth')
        torch.save(ema_model.module.state_dict(), f'{MODEL_DIR}ema_model_epoch{epoch}.pth')
        torch.save(optimizer.state_dict(), f'{MODEL_DIR}optimizer_epoch{epoch}.pth')
        torch.save(scaler.state_dict(), f'{MODEL_DIR}mp_scaler_epoch{epoch}.pth')
        torch.save(loss_fn.arcloss.state_dict(), f'{MODEL_DIR}arcloss_epoch{epoch}.pth')
        
        # end time of epoch
        epoch_end = time.perf_counter()
        print(f"epoch runtime: {epoch_end-epoch_start:5.3f} sec.")
        
        del logs_df, epoch_start, epoch_end
        torch.cuda.empty_cache()
                
    ################################## EMA Model Validation ################################  
    del model
    torch.cuda.empty_cache()
    
    ema_net = ema_model.module
    ema_net.eval()
    
    with torch.no_grad():
        # iterate over validation batches
        for (inputs, labels, ccm, meta) in valid_loader:
            inputs = inputs.to(device, non_blocking=True)
            meta = meta.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            ccm = ccm.to(device, non_blocking=True)

            # forward with mixed precision
            with autocast(device_type='cuda', dtype=torch.float16):                   
                outputs, embeddings = model(inputs, meta) 
                loss = loss_fn(outputs, embeddings, labels)

            # Compute metrics
            loss_metric.update(loss.detach())

            preds = outputs.softmax(dim=-1).detach()
            metrics.update(preds, labels)
            metric_ccm.update(preds * ccm, labels)

    # compute, sync & reset metrics for validation
    epoch_loss = loss_metric.compute()
    epoch_metrics = metrics.compute()
    epoch_metric_ccm = metric_ccm.compute()

    loss_metric.reset()
    metrics.reset()
    metric_ccm.reset()

    print(f"ema_loss: {epoch_loss.cpu().item():.5f}, ema_acc: {epoch_metrics['acc'].cpu().item():.5f}, ema_acc_top3: {epoch_metrics['top3_acc'].cpu().item():.5f}, ema_f1: {epoch_metrics['f1'].cpu().item():.5f}, ema_f1country: {epoch_metric_ccm.detach().cpu().item():.5f}") 
    
    with open(f'{MODEL_DIR}ema_results.txt', 'w') as f:
        print(f"ema_loss: {epoch_loss.cpu().item():.5f}, ema_acc: {epoch_metrics['acc'].cpu().item():.5f}, ema_acc_top3: {epoch_metrics['top3_acc'].cpu().item():.5f}, ema_f1: {epoch_metrics['f1'].cpu().item():.5f}, ema_f1country: {epoch_metric_ccm.detach().cpu().item():.5f}", file=f) 

    plot_history(logs)
    # end time of trainig
    end_training = time.perf_counter()
    print(f'Training succeeded in {(end_training - start_training):5.3f}s')

    if WANDB:
        wandb.finish()


if __name__=="__main__":
    main()