train.py

"""
Project: CVS-AdaptNet
-----
Copyright (c) University of Strasbourg, All Rights Reserved.
"""

import torch.nn as nn
import argparse
import torch
import clip
import torch.distributed as dist
from PIL import Image
from collections import defaultdict
from mmengine.logging import MMLogger
from torchmetrics import AveragePrecision as AP, Precision, Recall, F1Score
from torch import Tensor
import sys
import json

from codes.datasets import build_dataset
from codes.models import build_algorithm
from mmengine.config import Config
from transformers import AutoTokenizer
import torchmetrics
import numpy as np
import gc
from torch.utils.data import ConcatDataset
import torch.optim as optim
import torchvision.transforms as transforms
from tqdm import tqdm 
from codes.models.algorithms.CVS_AdaptNet import CVSAdaptNetModel, process_text
import json

import shutil
import os
import sys



def save_script(destination_folder, extra_files=None):
    # Get the path of the currently running script
    script_path = os.path.abspath(sys.argv[0])
    script_filename = os.path.basename(script_path)

    # Make sure the destination folder exists
    os.makedirs(destination_folder, exist_ok=True)

    # Save the training script itself
    destination_script_path = os.path.join(destination_folder, script_filename)
    shutil.copy(script_path, destination_script_path)

    # Save extra files (configs, prompts, etc.)
    if extra_files:
        for f in extra_files:
            if os.path.exists(f):
                shutil.copy(f, os.path.join(destination_folder, os.path.basename(f)))
            else:
                print(f"Warning: file not found: {f}")


def test_both(val_loader, model_1, args):
    class_prompt=args.class_prompt

    model_1.eval()

    with open(class_prompt) as f:
        lines = f.readlines()
    f.close()

    class_texts_ = [i.replace('\n', '') for i in lines]
    class_texts = process_text(class_texts_)
    text_features = model_1.model(None, class_texts, mode='text')['text_emb'].cuda()
    text_features /= text_features.norm(dim=-1, keepdim=True)
        # init logger
    logger: MMLogger = MMLogger.get_current_instance()
    logger_info = []
    eval_results = {}
    probs_list = []
    label_list = []
    with torch.no_grad():
        for val_loader in val_loaders:
            for i, data in enumerate(val_loader): 
                frames = data['video'].cuda() # (1, M, T, C, H, W)
                B, C, H, W = frames.shape

                frames = frames.view(-1, C, H, W)
                image_features = model_1.model(frames, None, mode='video')['img_emb'] # (B*M*T, D)

                image_features /= image_features.norm(dim=-1, keepdim=True)

                probs = (image_features @ text_features.T) # (1, classes)

                ## for average precision calculation
                preds = probs.cpu()
                
                gts = [[float(item) for item in json.loads(value)] for value in data['label']]


                probs_list = probs_list + [p.tolist() for p in preds]
                #gt_prob_list = label_list + [g for g in gts]
                label_list = label_list + [g for g in gts]

        torch_ap = AP(task='multilabel', num_labels=3, average='none')
        ds_preds = torch.stack([Tensor(p) for p in probs_list]).sigmoid()
        ds_gt = torch.stack([Tensor(g).round() for g in label_list]).long()

        ds_ap = torch_ap(ds_preds, ds_gt)
        #print(ds_ap)

        # log overall
        logger_info.append(f'ds_average_precision: {torch.nanmean(ds_ap):.4f}')
        eval_results['ds_average_precision'] = torch.nanmean(ds_ap)

        # log component-wise
        for ind, i in enumerate(ds_ap):
            logger_info.append(f'ds_average_precision_C{ind+1}: {i:.4f}')
            eval_results['ds_average_precision_C{}'.format(ind+1)] = i

        print('\n'.join(logger_info))
    return ds_ap

def test(classifier, val_loader, model_1, args):
    class_prompt=args.class_prompt

    model_1.eval()
    logger: MMLogger = MMLogger.get_current_instance()
    logger_info = []
    eval_results = {}
    probs_list = []
    label_list = []
    with torch.no_grad():
        for val_loader in val_loaders:
            for i, data in enumerate(val_loader): 
                frames = data['video'].cuda() # (1, M, T, C, H, W)
                # B, M, T, C, H, W = frames.shape
                B, C, H, W = frames.shape

                frames = frames.view(-1, C, H, W)
                image_features = model_1.model(frames)['img_emb']


                probs = nn.Sigmoid()(classifier(image_features))
                #print(probs)

                ## for average precision calculation
                preds = probs.cpu()
                
                gts = [[float(item) for item in json.loads(value)] for value in data['label']]
                probs_list = probs_list + [p.tolist() for p in preds]
                label_list = label_list + [g for g in gts]


        torch_ap = AP(task='multilabel', num_labels=3, average='none')
        ds_preds = torch.stack([Tensor(p) for p in probs_list])
        ds_gt = torch.stack([Tensor(g).round() for g in label_list]).long()
        ds_ap = torch_ap(ds_preds, ds_gt.long())
        #print(ds_ap)
        
        # log overall
        logger_info.append(f'ds_average_precision: {torch.nanmean(ds_ap):.4f}')
        eval_results['ds_average_precision'] = torch.nanmean(ds_ap)


        # log component-wise
        for ind, i in enumerate(ds_ap):
            logger_info.append(f'ds_average_precision_C{ind+1}: {i:.4f}')
            eval_results['ds_average_precision_C{}'.format(ind+1)] = i

        print('\n'.join(logger_info))
        return ds_ap



def fine_tune_image(
    train_loader: torch.utils.data.DataLoader,
    val_loader: torch.utils.data.DataLoader,
    model: torch.nn.Module,
    num_classes: int
) -> torch.nn.Module:

    
    model_clip = CVSAdaptNetModel(args, configs).cuda()
    model_clip.model = model

    for param in model_clip.model.parameters():
        param.requires_grad = True

        
    optimizer = torch.optim.Adam([
        {'params': model_clip.model.parameters(), 'lr': 0.00001}, # Fine-tune the pre-trained model
    ], weight_decay=0.0005)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=40)
    best_map = -1
    num_epochs = 20
    #epoch = 0
    best_map = -1
    json_write = {}
    for epoch in range(num_epochs):
        #epoch+=1
        model_clip.train()
        for batch in tqdm(train_loader):
            optimizer.zero_grad()
            
            classifier, total_loss = model_clip(batch,return_image=True)
            #Backward and optimize
            
            total_loss.backward()
            optimizer.step()
            print('epoch', epoch)
            print(total_loss.item())

        # Learning rate scheduling
        scheduler.step()
        print("_____________________________Before Cache clearance_________________________")
        print(torch.cuda.memory_summary(device=None, abbreviated=False))
        torch.cuda.empty_cache()
        print("_____________________________After Cache clearance_________________________")
        print(torch.cuda.memory_summary(device=None, abbreviated=False))
        gc.collect()
        with torch.no_grad():
            ds_ap = test(classifier, val_loader, model_clip, args)
            json_write[epoch] = str(ds_ap.detach().cpu().numpy())
            with open(args.save_dir +'/ds_ap_test_updated.json', 'w', encoding ='utf8') as json_file: 
                json.dump(json_write, json_file) 
                
            map_value = torch.nanmean(ds_ap)
            if map_value > best_map:
                # Update the best MAP value
                best_map = map_value
                #Save the model checkpoint
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model_clip.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'best_map': best_map
                }, args.save_dir +'/best_model_checkpoint_'+str(epoch)+'.pth')
                torch.save(model_clip, args.save_dir + "/best_model_checkpoint_"+str(epoch)+".pkl")
        
    return model_clip, classifier # Return the fine-tuned model and classifier


def fine_tune(
    train_loader: torch.utils.data.DataLoader,
    val_loader: torch.utils.data.DataLoader,
    model: torch.nn.Module,
    num_classes: int
) -> torch.nn.Module:

    model_clip = CVSAdaptNetModel(args, configs).cuda()
    model_clip.model = model
        # Freeze the pre-trained model's text parameters
    for param in model_clip.model.parameters():
        param.requires_grad = True
        
    optimizer = torch.optim.Adam([
        {'params': model_clip.model.parameters(), 'lr': 0.00001}, # Fine-tune the pre-trained model
        {'params': model_clip.logit_scale, 'lr': 0.001}, # logit scale is learnable
    ], weight_decay=0.0005)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=40)
    best_map = -1
    num_epochs = 20
    #epoch = 0
    ap_epoch = {}
    
    for epoch in range(num_epochs):
        model_clip.train()
        for batch in tqdm(train_loader):
            optimizer.zero_grad()  # Clear gradients
            total_loss = model_clip(batch,return_image=False)
            total_loss.backward()
            optimizer.step()
            print('epoch', epoch)
            print(total_loss.item())

        # Learning rate scheduling
        scheduler.step()
        print("_____________________________Before Cache clearance_________________________")
        print(torch.cuda.memory_summary(device=None, abbreviated=False))
        torch.cuda.empty_cache()
        print("_____________________________After Cache clearance_________________________")
        print(torch.cuda.memory_summary(device=None, abbreviated=False))
        gc.collect()

        # # Validation using both text and image
        with torch.no_grad():
            
            ds_ap = test_both(val_loader, model_clip, args)
            ap_epoch[epoch] = str(ds_ap.detach().cpu().numpy())
            with open(args.save_dir + '/ds_ap_both.json', 'w', encoding ='utf8') as json_file_both: 
                json.dump(ap_epoch, json_file_both) 
                
            map_value = torch.nanmean(ds_ap)
            if map_value > best_map:
                # Update the best MAP value
                best_map = map_value

                #Save the model checkpoint
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model_clip.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'best_map': best_map
                }, args.save_dir +'/best_model_checkpoint_both.pth')
                torch.save(model_clip, args.save_dir + "/best_model_checkpoint_both.pkl")

    return model_clip 
            


def get_args(description='CVSAdaptNet'):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--class_prompt', default='../class_prompt.txt', type=str, help='prompt for categories')
    parser.add_argument('--cvs_config', default='config/config_cvsadaptnet.py', type=str, help='config for dataset and cvsadaptnet')
    parser.add_argument('--checkpoint', default='./checkpoint.pkl', type=str, help='checkpoint of pre-trained weights')
    parser.add_argument('--batch_size', default=1, type=int, help='batch for training')
    parser.add_argument('--data_mode', default="both", type=str, help='mode for finetuning')
    parser.add_argument('--save_dir', default="save", type=str, help='directory for saving weights')
    parser.add_argument('--num_classes', default="3", type=int, help='number of classes')
    args = parser.parse_args()
    return args, parser



if __name__ == "__main__":
    
    args, _ = get_args()
    save_script(args.save_dir, extra_files=[args.cvs_config, args.class_prompt])

    device = "cuda" if torch.cuda.is_available() else "cpu"
    configs = Config.fromfile(args.cvs_config)['config']
    num_classes = args.num_classes
    model = build_algorithm(configs.model_config).cuda()
    print("Configuration loaded.")

    # Load model
    
    state_dict = torch.load(args.checkpoint)
    a, b = model.load_state_dict(state_dict, strict=True)


    train_datasets = [build_dataset(c) for c in configs.train_config]
    train_dataset = ConcatDataset(train_datasets)

    val_datasets = [build_dataset(c) for c in configs.val_config]
    val_dataset = ConcatDataset(val_datasets)


    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=4,
        pin_memory=True,

    )

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=4
    )

    val_loaders = [torch.utils.data.DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=0
    ) for val_dataset in val_datasets]

    print(args)

    
    if args.data_mode == "both":
        model_1 = fine_tune(train_loader, val_loader, model, num_classes)
    elif args.data_mode == "image":
        classifier, model_1 = fine_tune_image(train_loader, val_loader, model, num_classes)