trainer.py

import torch 
import gc 
import os

import pandas as pd
from tqdm import tqdm

from utils.metric import get_rd


def train(args, model, train_dataloader, test_dataloader, optimizer, scheduler, device):
    losses = []
    epohcs = args.epochs
    best_test_loss = float('inf')
    best_epoch = 0
    
    for epoch in tqdm(range(epohcs)):
        torch.cuda.empty_cache()
        gc.collect()
        
        model.train()
                
        train_loss_list = []
        test_loss_list = []
        print("Epoch:", epoch)

        for idx, batch in enumerate(train_dataloader):
            torch.cuda.empty_cache()
            gc.collect()
            labels = batch.pop("labels").to(device)
            flattened_patches = batch.pop("flattened_patches").to(device)
            attention_mask = batch.pop("attention_mask").to(device)
            
            if args.phase == 1:
                loss = model(flattened_patches = flattened_patches, 
                             attention_mask = attention_mask,
                             labels=labels)
            else:
                loss = model.forward_phase_2(flattened_patches = flattened_patches, 
                                             attention_mask = attention_mask,
                                             labels=labels,
                                             batch_size = args.batch_size)

            print("Loss:", loss.item())
            train_loss_list.append(loss.item())

            loss.backward()

            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()
            
        train_loss = sum(train_loss_list) / len(train_dataloader)
        
        del loss  
        del flattened_patches
        del attention_mask
        del labels
                    
        torch.cuda.empty_cache()
        gc.collect()
                  
        model.eval() 
        
        for idx, batch in enumerate(test_dataloader):
            labels = batch.pop("labels").to(device)
            flattened_patches = batch.pop("flattened_patches").to(device)
            attention_mask = batch.pop("attention_mask").to(device)
            
            loss = model(flattened_patches = flattened_patches, 
                            attention_mask = attention_mask,
                            labels=labels)
            test_loss_list.append(loss.item())
            
        del loss
        del flattened_patches
        del attention_mask
        del labels
                
        test_loss = sum(test_loss_list) / len(test_dataloader)
        losses.append({'epoch' : epoch, 'train_loss': train_loss ,'test_loss' : test_loss})
        
        if test_loss < best_test_loss:
            best_test_loss = test_loss
            best_epoch = epoch
            torch.save({'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'scheduler_state_dict': scheduler.state_dict()},
                       f'{args.model_save_path}/phase_{args.phase}_best_model.pth')
        
        summary = pd.DataFrame(losses)
        summary.to_csv(f'{args.model_save_path}/summary.csv')
        
        torch.save({'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict()},
                   f'{args.model_save_path}/phase_{args.phase}_{epoch}_model_state_dict.pth')

    print(f"Best model saved with test loss {best_test_loss} at epoch {best_epoch}")

def inference(args, model, dataset, dataloader, processor, device):
    os.makedirs(args.result_path, exist_ok=True)  # Ensure result path exists
    
    if args.inference_type == 'QA':
        accuracy_list = {'img': [], 'type': [], 'pred': [], 'label': []}
        for idx, batch in enumerate(dataloader):
            try:
                print(f"Processing batch {idx + 1}/{len(dataloader)}")
                labels = batch.pop("labels").to(device)
                flattened_patches = batch.pop("flattened_patches").to(device)
                attention_mask = batch.pop("attention_mask").to(device)
                chart_type = batch.pop("type")

                predictions = model.generate(flattened_patches=flattened_patches,
                                              attention_mask=attention_mask)
                pred = processor.batch_decode(predictions, skip_special_tokens=True)
                label = processor.batch_decode(labels, skip_special_tokens=True)

                accuracy_list['img'].append(dataset[idx]['img_name'])
                accuracy_list['type'].append(chart_type[0])
                accuracy_list['pred'].append(pred[0])
                accuracy_list['label'].append(label[0])
            except Exception as e:
                print(f"Error in batch {idx + 1}: {e}")
                continue  # Skip problematic batches

        result_df = pd.DataFrame(accuracy_list)
        print("Saving QA results...")
        result_df.to_csv(os.path.join(args.result_path, 'prediction.csv'), index=False)

    else:
        accuracy_list = {'img': [], 'pred': []}
        for idx, batch in enumerate(dataloader):
            try:
                print(f"Processing batch {idx + 1}/{len(dataloader)}")
                flattened_patches = batch.pop("flattened_patches").to(device)
                attention_mask = batch.pop("attention_mask").to(device)

                predictions = model.generate(flattened_patches=flattened_patches,
                                              attention_mask=attention_mask)
                pred = processor.batch_decode(predictions, skip_special_tokens=True)

                accuracy_list['img'].append(dataset[idx]['img_name'])
                accuracy_list['pred'].append(pred[0])
            except Exception as e:
                print(f"Error in batch {idx + 1}: {e}")
                continue

        result_df = pd.DataFrame(accuracy_list)
        print("Saving OpenCQA results...")
        result_df.to_csv(os.path.join(args.result_path, 'opencqa_prediction.csv'), index=False)

    print("Inference complete.")


# def inference(args, model, dataset, dataloader, processor, device):
    
#     if args.inference_type == 'QA':
#         accuracy_list = {'img':[], 'type':[], 'pred':[], 'label':[]}

#         for idx, batch in enumerate(dataloader):
#             labels = batch.pop("labels").to(device)
#             flattened_patches = batch.pop("flattened_patches").to(device)
#             attention_mask = batch.pop("attention_mask").to(device)
#             chart_type = batch.pop("type")
            
#             predictions = model.generate(flattened_patches= flattened_patches,
#                                         attention_mask=attention_mask,)    
#             pred = processor.batch_decode(predictions, skip_special_tokens=True)
#             label = processor.batch_decode(labels, skip_special_tokens=True)
            
#             accuracy_list['img'].append(dataset[idx]['img_name'])
#             accuracy_list['type'].append(chart_type[0])
#             accuracy_list['pred'].append(pred[0])
#             accuracy_list['label'].append(label[0])

#         result_df = pd.DataFrame(accuracy_list)
        
        
#         rd_df, failed = get_rd(result_df)
#         rd_df.to_csv(os.path.join(args.result_path, 'prediction.csv'))
    
#     else:
#         accuracy_list = {'img':[], 'pred':[]}
#         for idx, batch in enumerate(dataloader):
#             flattened_patches = batch.pop("flattened_patches").to(device)
#             attention_mask = batch.pop("attention_mask").to(device)
            
#             predictions = model.generate(flattened_patches= flattened_patches,
#                                         attention_mask=attention_mask)    
#             pred = processor.batch_decode(predictions, skip_special_tokens=True)
            
#             accuracy_list['img'].append(dataset[idx]['img_name'])
#             accuracy_list['pred'].append(pred[0])
            
#         result_df = pd.DataFrame(accuracy_list)
#         result_df.to_csv(os.path.join(args.result_path, 'opencqa_prediction.csv'))