train_decoding.py

"""
Training script for the Dewave EEG-to-Text decoding model.
Implements a two-step training process with customizable configurations.
"""

import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
import pickle
import json
import time
import copy
from tqdm import tqdm
from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig
# from original_data import ZuCo_dataset
# from data import ZuCo_dataset
from data_for_rawEEG import ZuCo_dataset
# from new_data import ZuCo_dataset
from model_decoding import Dewave
from config import get_config
from types_ import *
from eval_decoding import eval_model

def train_model(
    dataloaders: Dict[str, DataLoader],
    device: torch.device,
    model: nn.Module,
    criterion: nn.Module,
    optimizer: optim.Optimizer,
    scheduler: lr_scheduler._LRScheduler,
    num_epochs: int = 25,
    checkpoint_path_best: str = './checkpoints/decoding/best/temp_decoding.pt',
    checkpoint_path_last: str = './checkpoints/decoding/last/temp_decoding.pt',
    stage: str = 'first'
) -> nn.Module:
    """
    Train the Dewave model with validation.

    Args:
        dataloaders: Dictionary containing train and dev dataloaders
        device: Device to run the training on
        model: Model to train
        criterion: Loss criterion
        optimizer: Optimizer for training
        scheduler: Learning rate scheduler
        num_epochs: Number of epochs to train
        checkpoint_path_best: Path to save best model checkpoint
        checkpoint_path_last: Path to save last model checkpoint
        stage: Training stage identifier

    Returns:
        nn.Module: Trained model with best weights
    """
    # modified from: https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
    since = time.time()
      
    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = 100000000000

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'dev']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0

            # Iterate over data.
            for sent_level_EEG, rawEEG, input_masks, input_mask_invert, target_ids, target_mask, real_length in tqdm(dataloaders[phase]):

                # load in batch
                input_embeddings_batch = rawEEG.to(device).float()
                # input_embeddings_batch = input_embeddings.to(device).float()
                input_masks_batch = input_masks.to(device)
                input_mask_invert_batch = input_mask_invert.to(device)
                # input_masks_batch = torch.ones([rawEEG.shape[0],rawEEG.shape[2]]).to(device)
                target_ids_batch = target_ids.to(device)
                context = target_ids_batch.clone()
                """replace padding ids in target_ids with -100"""
                target_ids_batch[target_ids_batch == tokenizer.pad_token_id] = -100

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
    	        # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    out, loss = model(input_embeddings_batch, context, input_masks_batch, input_mask_invert_batch, target_ids_batch, stage=stage)# out.logits shape 8*48*50265
                                 
                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.sum().backward()
                        # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                        optimizer.step()

                # statistics
                running_loss += loss.sum().item() * input_embeddings_batch.size()[0] # batch loss
                

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]

            print('{} Loss: {:.4f}'.format(phase, epoch_loss))

            # deep copy the model
            if phase == 'dev' and epoch_loss < best_loss:
                best_loss = epoch_loss
                best_model_wts = copy.deepcopy(model.state_dict())
                '''save checkpoint'''
                torch.save(model.state_dict(), checkpoint_path_best)
                print(f'update best on dev checkpoint: {checkpoint_path_best}')
        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val loss: {:4f}'.format(best_loss))
    torch.save(model.state_dict(), checkpoint_path_last)
    print(f'update last checkpoint: {checkpoint_path_last}')

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

def show_require_grad_layers(model):
    print()
    print(' require_grad layers:')
    # sanity check
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(' ', name)

if __name__ == '__main__':
    args = get_config('train_decoding')

    input_mode = args['train_input']
    
    ''' config param'''
    dataset_setting = 'unique_sent'
    
    num_epochs_step1 = args['num_epoch_step1']
    num_epochs_step2 = args['num_epoch_step2']
    step1_lr = args['learning_rate_step1']
    step2_lr = args['learning_rate_step2']
    
    batch_size = args['batch_size']
    
    model_name = args['model_name']

    save_name = args['save_name']

    # task_name = 'task1'
    # task_name = 'task1_task2'
    # task_name = 'task1_task2_task3'
    # task_name = 'task1_task2_taskNRv2'
    task_name = args['task_name']
    train_input = args['train_input']
    print("train_input is:", train_input)   
    save_path = args['save_path']
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    skip_step_one = args['skip_step_one']
    load_step1_checkpoint = args['load_step1_checkpoint']
    use_random_init = args['use_random_init']
    device_ids = [0,1] # device setting

    if use_random_init and skip_step_one:
        step2_lr = 5*1e-4
        
    print(f'[INFO]using model: {model_name}')
    dataset_setting = args['dataset_setting']


    # if skip_step_one:
    #     save_name = f'{task_name}_finetune_{model_name}_skipstep1_b{batch_size}_{num_epochs_step1}_{num_epochs_step2}_{step1_lr}_{step2_lr}_{dataset_setting}_{train_input}'
    # else:
    #     save_name = f'{task_name}_finetune_{model_name}_2steptraining_b{batch_size}_{num_epochs_step1}_{num_epochs_step2}_{step1_lr}_{step2_lr}_{dataset_setting}_{train_input}'
    
    # if use_random_init:
    #     save_name = 'randinit_' + save_name
    # save_name= dataset_setting
    save_path_best = os.path.join(save_path, 'best')
    if not os.path.exists(save_path_best):
        os.makedirs(save_path_best)

    output_checkpoint_name_best = os.path.join(save_path_best, f'{save_name}.pt')

    save_path_last = os.path.join(save_path, 'last')
    if not os.path.exists(save_path_last):
        os.makedirs(save_path_last)

    output_checkpoint_name_last = os.path.join(save_path_last, f'{save_name}.pt')

    # subject_choice = 'ALL
    subject_choice = args['subjects']
    print(f'![Debug]using {subject_choice}')
    # eeg_type_choice = 'GD
    eeg_type_choice = args['eeg_type']
    print(f'[INFO]eeg type {eeg_type_choice}')
    # bands_choice = ['_t1','_t2','_a1','_a2','_b1','_b2','_g1','_g2'] 
    bands_choice = args['eeg_bands']
    print(f'[INFO]using bands {bands_choice}')


    
    ''' set random seeds '''
    seed_val = 312
    # random.seed(seed_val)
    os.environ['PYTHONHASHSEED'] = str(seed_val)
    np.random.seed(seed_val)
    torch.manual_seed(seed_val)
    torch.cuda.manual_seed_all(seed_val)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    ''' set up device '''
    # use cuda
    if torch.cuda.is_available():  
        # dev = "cuda:3" 
        dev = args['cuda'] 
    else:  
        dev = "cpu"
    # CUDA_VISIBLE_DEVICES=0,1,2,3  
    device = torch.device(dev)
    print(f'[INFO]using device {dev}')
    print()
    ''' set up dataloader '''
    
    whole_dataset_dicts = []

    if 'task1' in task_name:
        dataset_path_task1 = './datasets/rawZuCo/task1-SR/task1-SR-dataset.pickle'
        with open(dataset_path_task1, 'rb') as handle:
            whole_dataset_dicts.append(pickle.load(handle))
    if 'task2' in task_name:
        dataset_path_task2 = './datasets/rawZuCo/task2-NR/task2-NR-dataset.pickle'
        with open(dataset_path_task2, 'rb') as handle:
            whole_dataset_dicts.append(pickle.load(handle))
    if 'task3' in task_name:
        dataset_path_task3 = './datasets/rawZuCo/task3-TSR/task3-TSR-dataset.pickle'
        with open(dataset_path_task3, 'rb') as handle:
            whole_dataset_dicts.append(pickle.load(handle))
    if 'taskNRv2' in task_name:
        dataset_path_taskNRv2 = './datasets/rawZuCo/task2-NR-2.0/task2-NR-2.0-dataset.pickle'
        with open(dataset_path_taskNRv2, 'rb') as handle:
            whole_dataset_dicts.append(pickle.load(handle))


    """save config"""
    cfg_dir = './config/decoding/'

    if not os.path.exists(cfg_dir):
        os.makedirs(cfg_dir)

    with open(os.path.join(cfg_dir,f'{save_name}.json'), 'w') as out_config:
        json.dump(args, out_config, indent = 4)


    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
    # subject ['ZAB', 'ZDM', 'ZGW', 'ZJM', 'ZJN', 'ZJS', 'ZKB', 'ZKH', 'ZKW', 'ZMG', 'ZPH']

    if dataset_setting != 'unique_sent':


        subjects=['YDR', 'YFR', 'YFS', 'YAC', 'YDG', 'YHS', 'YMD', 'YLS', 'YRH', 'YRK', 'YSD', 'YRP', 'YSL', 'YTL', 'YAG', 'YAK', 'YIS', 'YMS']
        
        for subj in subjects:

            # train dataset
            train_set = ZuCo_dataset(whole_dataset_dicts, 'train', tokenizer, eeg_type = eeg_type_choice, bands = bands_choice, setting = subj)
            # dev dataset
            dev_set = ZuCo_dataset(whole_dataset_dicts, 'dev', tokenizer, eeg_type = eeg_type_choice, bands = bands_choice, setting = subj)

            
            dataset_sizes = {'train': len(train_set), 'dev': len(dev_set)}
            print('[INFO]train_set size: ', len(train_set))
            print('[INFO]dev_set size: ', len(dev_set))
            
            
            # train dataloader
            train_dataloader = DataLoader(train_set, batch_size = batch_size, shuffle=True, num_workers=4)
            # dev dataloader
            val_dataloader = DataLoader(dev_set, batch_size = 1, shuffle=False, num_workers=4)
            # test dataloader
            
            # dataloaders
            dataloaders = {'train':train_dataloader, 'dev':val_dataloader}

            # ''' set up model '''
            # if model_name == 'BrainTranslator':
            #     if use_random_init:
            #         config = BartConfig.from_pretrained('facebook/bart-large')
            #         pretrained = BartForConditionalGeneration(config)
            #     else:
            pretrained = BartForConditionalGeneration.from_pretrained('facebook/bart-large')

            model = Dewave(pretrained, input_type=input_mode, embedding_dim = 2048, num_embeddings = 512)
            # model = Dewave(pretrained, input_type='EEG', embedding_dim = 2048, num_embeddings = 512)
            model.to(device)
            # model = torch.nn.DataParallel(model, device_ids=device_ids)
            
            ''' training loop '''

            ######################################################
            '''step one trainig: freeze most of BART params'''
            ######################################################

            # closely follow BART paper

            for name, param in model.named_parameters():
                if param.requires_grad and 'pretrained' in name:
                    if ('shared' in name) or ('embed_positions' in name) or ('encoder.layers.0' in name):
                        continue
                    else:
                        param.requires_grad = False

            ''' set up loss function '''
            criterion = nn.CrossEntropyLoss()


            # SGD algorithm with a learning rate of 0.0005 for 30 epochs with 0.1 times the learning rate decrease at epoch 20

            if input_mode == 'rawEEG': 
                optimizer_step_pretrian = optim.SGD(model.parameters(), lr=0.0005, momentum=0.9)

                exp_lr_scheduler_pretrian = lr_scheduler.StepLR(optimizer_step_pretrian, step_size=30, gamma=0.1)
                model = train_model(dataloaders, device, model, criterion, optimizer_step_pretrian, exp_lr_scheduler_pretrian, num_epochs=30, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'pretrained')

            if skip_step_one:
                if load_step1_checkpoint:
                    stepone_checkpoint = 'path_to_step_1_checkpoint.pt'
                    print(f'skip step one, load checkpoint: {stepone_checkpoint}')
                    model.load_state_dict(torch.load(stepone_checkpoint))
                else:
                    print('skip step one, start from scratch at step two')
            else:
                ''' set up optimizer and scheduler'''
                optimizer_step1 = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=step1_lr, momentum=0.9)

                exp_lr_scheduler_step1 = lr_scheduler.StepLR(optimizer_step1, step_size=20, gamma=0.1)

                print('=== start Step1 training ... ===')
                # print training layers
                # show_require_grad_layers(model)
                # return best loss model from step1 training
                # optimized by the SGD optimizer with a learning rate of 0.0005 for 35 epochs
                model = train_model(dataloaders, device, model, criterion, optimizer_step1, exp_lr_scheduler_step1, num_epochs=num_epochs_step1, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'first')

            ######################################################
            '''step two trainig: update whole model for a few iterations'''
            ######################################################

            # language model θBART is opened to fine-tune the whole system.
            for name, param in model.named_parameters():
                param.requires_grad = True

            ''' set up optimizer and scheduler'''
            optimizer_step2 = optim.SGD(model.parameters(), lr=step2_lr, momentum=0.9)

            exp_lr_scheduler_step2 = lr_scheduler.StepLR(optimizer_step2, step_size=30, gamma=0.1)

            print()
            print('=== start Step2 training ... ===')

            # print training layers
            # show_require_grad_layers(model)
            # test dataset


            '''second stage'''
            # the second stage is optimized by the SGD optimizer with a learning rate of 1e − 6 for 35 epochs.
            trained_model = train_model(dataloaders, device, model, criterion, optimizer_step2, exp_lr_scheduler_step2, num_epochs=num_epochs_step2, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'second')

            test_set = ZuCo_dataset(whole_dataset_dicts, 'test', tokenizer, eeg_type = eeg_type_choice, bands = bands_choice, setting = subj)
            print('[INFO]test_set size: ', len(test_set))
            test_dataloader = DataLoader(test_set, batch_size = 1, shuffle=False, num_workers=4)
            dataloaders = {'test':test_dataloader}
            eval_model(dataloaders, device, tokenizer, model=trained_model, output_all_results_path=f'./results/{dataset_setting}-{subj}.txt', score_results = f'./score_results/{dataset_setting}-{subj}.txt')
            del model, optimizer_step2, exp_lr_scheduler_step2
            torch.cuda.empty_cache()

    else:

        dataset_setting = 'unique_sent'
        subject_choice = 'ALL'
        train_set = ZuCo_dataset(whole_dataset_dicts, 'train', tokenizer, subject = subject_choice, bands = bands_choice, setting = dataset_setting)
        # dev dataset
        dev_set = ZuCo_dataset(whole_dataset_dicts, 'dev', tokenizer, subject = subject_choice, bands = bands_choice, setting = dataset_setting)
            
        dataset_sizes = {'train': len(train_set), 'dev': len(dev_set)}
        print('[INFO]train_set size: ', len(train_set))
        print('[INFO]dev_set size: ', len(dev_set))
        
        
        # train dataloader
        train_dataloader = DataLoader(train_set, batch_size = batch_size, shuffle=True, num_workers=4)
        # dev dataloader
        val_dataloader = DataLoader(dev_set, batch_size = 1, shuffle=False, num_workers=4)
        # test dataloader
        
        # dataloaders
        dataloaders = {'train':train_dataloader, 'dev':val_dataloader}

        # ''' set up model '''
        # if model_name == 'BrainTranslator':
        #     if use_random_init:
        #         config = BartConfig.from_pretrained('facebook/bart-large')
        #         pretrained = BartForConditionalGeneration(config)
        #     else:
        pretrained = BartForConditionalGeneration.from_pretrained('facebook/bart-large')

        model = Dewave(pretrained, input_type=input_mode, embedding_dim = 2048, num_embeddings = 512)
        # model = Dewave(pretrained, input_type='EEG', embedding_dim = 2048, num_embeddings = 512)
        model.to(device)
        # model = torch.nn.DataParallel(model, device_ids=device_ids)
        
        ''' training loop '''

        ######################################################
        '''step one trainig: freeze most of BART params'''
        ######################################################

        # closely follow BART paper

        for name, param in model.named_parameters():
            if param.requires_grad and 'pretrained' in name:
                if ('shared' in name) or ('embed_positions' in name) or ('encoder.layers.0' in name):
                    continue
                else:
                    param.requires_grad = False

        ''' set up loss function '''
        criterion = nn.CrossEntropyLoss()


        # SGD algorithm with a learning rate of 0.0005 for 30 epochs with 0.1 times the learning rate decrease at epoch 20

        if input_mode == 'rawEEG': 
            optimizer_step_pretrian = optim.SGD(model.parameters(), lr=0.0005, momentum=0.9)

            exp_lr_scheduler_pretrian = lr_scheduler.StepLR(optimizer_step_pretrian, step_size=30, gamma=0.1)
            model = train_model(dataloaders, device, model, criterion, optimizer_step_pretrian, exp_lr_scheduler_pretrian, num_epochs=30, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'pretrained')

        if skip_step_one:
            if load_step1_checkpoint:
                stepone_checkpoint = 'path_to_step_1_checkpoint.pt'
                print(f'skip step one, load checkpoint: {stepone_checkpoint}')
                model.load_state_dict(torch.load(stepone_checkpoint))
            else:
                print('skip step one, start from scratch at step two')
        else:
            ''' set up optimizer and scheduler'''
            optimizer_step1 = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=step1_lr, momentum=0.9)

            exp_lr_scheduler_step1 = lr_scheduler.StepLR(optimizer_step1, step_size=20, gamma=0.1)

            print('=== start Step1 training ... ===')
            # print training layers
            # show_require_grad_layers(model)
            # return best loss model from step1 training
            # optimized by the SGD optimizer with a learning rate of 0.0005 for 35 epochs
            model = train_model(dataloaders, device, model, criterion, optimizer_step1, exp_lr_scheduler_step1, num_epochs=num_epochs_step1, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'first')

        ######################################################
        '''step two trainig: update whole model for a few iterations'''
        ######################################################

        # language model θBART is opened to fine-tune the whole system.
        for name, param in model.named_parameters():
            param.requires_grad = True

        ''' set up optimizer and scheduler'''
        optimizer_step2 = optim.SGD(model.parameters(), lr=step2_lr, momentum=0.9)

        exp_lr_scheduler_step2 = lr_scheduler.StepLR(optimizer_step2, step_size=30, gamma=0.1)

        print()
        print('=== start Step2 training ... ===')

        # print training layers
        # show_require_grad_layers(model)
        # test dataset


        '''second stage'''
        # the second stage is optimized by the SGD optimizer with a learning rate of 1e − 6 for 35 epochs.
        trained_model = train_model(dataloaders, device, model, criterion, optimizer_step2, exp_lr_scheduler_step2, num_epochs=num_epochs_step2, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last, stage = 'second')

        test_set = ZuCo_dataset(whole_dataset_dicts, 'test', tokenizer, subject = subject_choice, bands = bands_choice, setting = dataset_setting)
        print('[INFO]test_set size: ', len(test_set))
        test_dataloader = DataLoader(test_set, batch_size = 1, shuffle=False, num_workers=4)
        dataloaders = {'test':test_dataloader}
        eval_model(dataloaders, device, tokenizer, model=trained_model, output_all_results_path=f'./results/{dataset_setting}-{subject_choice}.txt', score_results = f'./score_results/{dataset_setting}-{subject_choice}.txt')
        del model, optimizer_step2, exp_lr_scheduler_step2
        torch.cuda.empty_cache()