main.py

import os
import copy
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from os.path import join
from omegaconf import OmegaConf

from dataloaders import initialize_data_functions, get_evaluation_loaders
from utils.logging import print_header, print_args, print_config
from optimizer import get_optimizer, get_scheduler
from loss import get_loss
from data_transforms import get_data_transforms
from train import train_model, evaluate_model, plot_forecasts

from setup import format_arg, seed_everything
from setup import initialize_args
from setup import load_model_config, load_main_config
from setup import initialize_experiment
from setup.configs.model import update_output_config_from_args  # For multivariate feature prediction

from model.network import SpaceTime
from customized_dataloader import read_partition_dataset, get_data_loader
import yaml
import pdb 
from accelerate import Accelerator
from utils.model_loader import model_loader
from model.attn import GPT
import torch.nn as nn
from accelerate.utils import InitProcessGroupKwargs
from datetime import timedelta

def load_config_to_args(args, config_file):
    with open(config_file, 'r') as f:
        config = yaml.safe_load(f)

    for key, value in config.items():
        if hasattr(args, key):
            setattr(args, key, value)

    return args

def main():
    print_header('*** EXPERIMENT ARGS ***')
    args = initialize_args()
    load_config_to_args(args, args.config_file)
    seed_everything(args.seed)
    experiment_configs = load_main_config(args, config_dir='./configs')
    
    load_data, visualize_data = initialize_data_functions(args)
    print_header('*** DATASET ***')
    print_config(experiment_configs['dataset'])
    
    print_header('*** LOADER ***')
    print_config(experiment_configs['loader'])
    
    print_header('*** OPTIMIZER ***')
    print_config(experiment_configs['optimizer'])
    
    print_header('*** SCHEDULER ***')
    print_config(experiment_configs['scheduler'])

    assert args.model in ('attn', 'attn_ts_v2', 'attn_lg_v2', 'spacetime', 'spacetime_v2', 'spacetime_lg',
        'spacetime_ts', 'spacetime_ts_lg', 'spacetime_lg_v2', 'spacetime_ts_v2', \
        'attn_pe', 'attn_pe_ts_v2', 'attn_pe_lg_v2', \
        'spacetime_med_ts_v2', 'spacetime_med_v2', 'attn_med_v2')
    
    if args.sampler:
        args.model = args.model + '_sampler'

    # Set timeout to 10 minutes
    timeout = timedelta(seconds=3600)
    process_group_kwargs = InitProcessGroupKwargs(timeout=timeout)

    # Pass kwargs_handlers to Accelerator
    accelerator = Accelerator(kwargs_handlers=[process_group_kwargs])
    # accelerator = Accelerator()
    device = accelerator.device

    # accelerator = None 
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Loading Data
    if args.dataset == 'customized':
        read_data = read_partition_dataset(
        # "./dataloaders/data/informer/etth/ETTh3.csv")
        'merged_csv_final.csv')
        dataloaders = get_data_loader(read_data,
                    lag=args.lag, horizon=args.horizon,
                    batch_size=args.batch_size, use_sampler=args.sampler)
    elif args.dataset == 'customized_a':
        read_data = read_partition_dataset(
        # "./dataloaders/data/informer/etth/ETTh3.csv")
        'all_domain_m2.csv')
        dataloaders = get_data_loader(read_data,
                    lag=args.lag, horizon=args.horizon,
                    batch_size=args.batch_size, use_sampler=args.sampler)
    elif args.dataset == 'customized_s':
        read_data = read_partition_dataset(
        "./dataloaders/data/informer/etth/ETTh3.csv")
        dataloaders = get_data_loader(read_data,
                    lag=args.lag, horizon=args.horizon,
                    batch_size=args.batch_size, use_sampler=args.sampler)
    else:
        dataloaders = load_data(experiment_configs['dataset'], 
                            experiment_configs['loader'])
        dataloaders = [d for d in dataloaders]
    
    train_loader, val_loader, test_loader = dataloaders
    
    splits = ['train', 'val', 'test']
    dataloaders_by_split = {split: dataloaders[ix] 
                            for ix, split in enumerate(splits)}

    eval_loaders = get_evaluation_loaders(dataloaders, batch_size=args.batch_size)
    
    # Setup input_dim based on features
    x, y, *z = train_loader.dataset.__getitem__(0)
    args.input_dim = x.shape[1]  # L x D
    output_dim = y.shape[1]
    
    # Initialize Model
    # args.device = (torch.device('cuda:0') 
    #                if torch.cuda.is_available() and not args.no_cuda
    #                else torch.device('cpu'))
    args.device = device 

    model_configs = {'embedding_config': args.embedding_config,
                     'encoder_config':   args.encoder_config,
                     'decoder_config':   args.decoder_config,
                     'output_config':    args.output_config}
    model_configs = OmegaConf.create(model_configs)
    model_configs = load_model_config(model_configs, config_dir='./configs/model',
                                      args=args)
    model_configs['inference_only'] = False
    model_configs['lag'] = args.lag
    model_configs['horizon'] = args.horizon
    
    if args.features == 'M':  # Update output
        update_output_config_from_args(model_configs['output_config'], args,
                                       update_output_dim=True, output_dim=output_dim)
        model_configs['output_config'].input_dim = model_configs['output_config'].kwargs.input_dim
        model_configs['output_config'].output_dim = model_configs['output_config'].kwargs.output_dim  
        print(model_configs['output_config'])

    use_ts, use_lg, use_pe = False, False, False
    if '_ts_v2' in args.model:
        use_ts = 'v2'
        model_configs['use_ts'] = use_ts
    elif '_ts' in args.model:
        use_ts = 'v1'
        model_configs['use_ts'] = use_ts

    if '_lg_v2' in args.model:
        use_lg = 'v2' 
        model_configs['use_lg'] = use_lg
    elif '_lg' in args.model:
        use_lg = 'v1' 
        model_configs['use_lg'] = use_lg
    
    if '_pe' in args.model:
        use_pe = True
        model_configs['use_pe'] = use_pe

    if 'attn' in args.model:
        model = GPT(lag=args.lag, horizon=args.horizon, d=args.kernel_dim, use_ts=use_ts, use_lg=use_lg, \
                    use_pe=use_pe)
    else:
        # if args.model == 'spacetime_ts':
        #     model_configs['use_ts'] = True
        # if args.model == 'spacetime_lg':
        #     model_configs['use_lg'] = True
        model = SpaceTime(**model_configs)
        model.replicate = args.replicate  # Only used for testing specific things indicated by replicate
        model.set_lag(args.lag)
        model.set_horizon(args.horizon)
    # Initialize optimizer and scheduler
    optimizer = get_optimizer(model, experiment_configs['optimizer'])
    scheduler = get_scheduler(model, optimizer, 
                              experiment_configs['scheduler'])
    # Save some model stats
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    args.model_parameter_count = params
    arg_dict = print_args(args, return_dict=True, verbose=args.verbose)
    # Setup logging
    wandb = initialize_experiment(args, experiment_name_id='',
                                  best_train_metric=1e10, 
                                  best_val_metric=1e10)
    try:
        pd.DataFrame.from_dict(arg_dict).to_csv(args.log_configs_path)
    except:
        pd.DataFrame.from_dict([arg_dict]).to_csv(args.log_configs_path)
        
    if args.verbose:
        print_header('*** MODEL ***')
        print(model)
        print_config(model_configs)
        
        from einops import rearrange
        _k = model.encoder.blocks[0].pre.get_kernel(rearrange(x, '(o l) d -> o d l', o=1))
        _k_diff = model.encoder.blocks[0].pre.diff_kernel
        _k_ma_r = model.encoder.blocks[0].pre.ma_r_kernel
        print_header(f'──> Preprocessing kernels (full: {_k.shape}, diff: {_k_diff.shape}, ma: {_k_ma_r.shape})')
        print(_k[:16, :_k_ma_r.shape[-1]])
                     
    
    print_header(f'*** TRAINING ***')
    print(f'├── Lag: {args.lag}')
    print(f'├── Horizon: {args.horizon}')
    print(f'├── Criterion: {args.loss}, weights: {args.criterion_weights}')
    print(f'├── Dims: input={args.input_dim}, model={args.model_dim}')
    print(f'├── Number trainable parameters: {params}')  # └── 
    print(f'├── Experiment name: {args.experiment_name}')
    print(f'├── Logging to: {args.log_results_path}')
    
    # Loss objectives
    criterions = {name: get_loss(name) for name in ['rmse', 'mse', 'mae', 'rse']}

    eval_criterions = criterions
    for name in ['rmse', 'mse', 'mae']:
        eval_criterions[f'informer_{name}'] = get_loss(f'informer_{name}')
    
    input_transform, output_transform = get_data_transforms(args.data_transform,
                                                            args.lag)
    
    # load existing checkpoint
    if args.exist_model is not None and \
        os.path.exists(args.exist_model):
        model = model_loader(model, args)

    # if torch.cuda.device_count() > 1:
    #     print(f"Using {torch.cuda.device_count()} GPUs!")
    #     model = nn.DataParallel(model)

    if accelerator is not None: 
        model, optimizer, \
            dataloaders[0], dataloaders[1], dataloaders[2] \
                = accelerator.prepare(model, optimizer, dataloaders[0], dataloaders[1], dataloaders[2])
    
    dataloaders_by_split = {split: dataloaders[ix] 
                            for ix, split in enumerate(splits)}

    model = train_model(model, optimizer, scheduler, dataloaders_by_split, 
                        criterions, max_epochs=args.max_epochs, config=args, 
                        input_transform=input_transform,
                        output_transform=output_transform,
                        val_metric=args.val_metric, wandb=wandb, 
                        return_best=True, early_stopping_epochs=args.early_stopping_epochs,
                        accelerator=accelerator)                     
    
if __name__ == '__main__':
    main()