show_gpu('Initial GPU memory usage:')
for i in range(2):
    model, optimizer, scheduler = get_training_obj(params)
    show_gpu(f'{i}: GPU memory usage after loading training objects:')
    for epoch in range(1):
        epoch_start = time.time()
        model.train()
        for batch in dp.train_dataloader:
            xb,mb,_,yb = tuple(t.to(params['device']) for t in batch)
            outputs = model(input_ids = xb, attention_mask = mb, labels = yb)
            loss = outputs[0]
            loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
    show_gpu(f'{i}: GPU memory usage after training model:')
    del model, optimizer, scheduler, loss, outputs
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    show_gpu(f'{i}: GPU memory usage after clearing cache:')

import platform; print("Platform", platform.platform())
import sys; print("Python", sys.version)
import torch; print("PyTorch", torch.__version__)

from __future__ import absolute_import, division, print_function
import glob
import logging
import os
import time
import json
import random
import numpy as np
import pandas as pd
from random import sample, seed
import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler,TensorDataset
from torch.utils.data.distributed import DistributedSampler
from transformers import RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer
from transformers import DistilBertConfig, DistilBertForSequenceClassification, DistilBertTokenizer
from transformers import BertConfig, BertForSequenceClassification, BertTokenizer
from transformers import AdamW, WarmupLinearSchedule
from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors
from transformers import glue_convert_examples_to_features as convert_examples_to_features
import subprocess

params = {
    'num_epochs': 2,
    'warmup_ratio': 0.06,
    'weight_decay': 0.1,
    'adam_epsilon': 1e-6,
    'model_name': 'roberta-large',
    'max_grad_norm': 1.0,
    'lr': 2e-5,
    'bs': 32,
    'device': 'cuda',
    'task': 'cola',
    'data_dir': '/home/ubuntu/glue_data/CoLA',
    'max_seq_length': 50,          
    'metric_name': 'mcc',
    'patience': 3,
    'seed': 935,
    'n': -1,
      }


class DataProcessor():
    '''Preprocess the data, store data loaders and tokenizer'''
    
    _TOKEN_TYPES = {
        'roberta': RobertaTokenizer,
        'distilbert': DistilBertTokenizer,
        'bert': BertTokenizer,
    }
    
    def __init__(self, params):
        model_type = params['model_name'].split('-')[0]
        assert model_type in self._TOKEN_TYPES.keys()
        self.tok = self._TOKEN_TYPES[model_type]
        self.params = params
        self.processor = processors[self.params['task']]()
        self.output_mode = output_modes[self.params['task']]
        self.label_list = self.processor.get_labels()
    
    @staticmethod
    def _convert_to_tensors(features):
        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
        all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
        all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
        return TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels)
    
    def _load_examples(self, tokenizer, evaluate):
        if evaluate:
            examples = self.processor.get_dev_examples(self.params['data_dir']) 
        else: 
            examples = self.processor.get_train_examples(self.params['data_dir'])
            if self.params['n'] >= 0:
                examples = sample(examples, self.params['n'])
        features = convert_examples_to_features(examples,
                                                tokenizer,
                                                label_list=self.label_list,
                                                max_length=self.params['max_seq_length'],
                                                output_mode=self.output_mode,
                                                pad_on_left=False,
                                                pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
                                                pad_token_segment_id=0)
        
        return self._convert_to_tensors(features)
    
    def _define_tokenizer(self):
        return self.tok.from_pretrained(self.params['model_name'], do_lower_case=True)
    
    def load_data(self):
        tokenizer = self._define_tokenizer()
        self.train_data = self._load_examples(tokenizer, False)
        self.valid_data = self._load_examples(tokenizer, True)
        self.train_n = len(self.train_data)
        self.valid_n = len(self.valid_data)
        self.params['total_steps'] = self.params['num_epochs'] * self.train_n
        return self.params
    
    def create_loaders(self):
        self.train_dataloader = DataLoader(self.train_data, shuffle=True, batch_size=self.params['bs'])
        self.valid_dataloader = DataLoader(self.valid_data, shuffle=False, batch_size=2*self.params['bs'])

dp = DataProcessor(params)
params = dp.load_data()
dp.create_loaders()


def show_gpu(msg):
    """
    ref: https://discuss.pytorch.org/t/access-gpu-memory-usage-in-pytorch/3192/4
    """
    def query(field):
        return(subprocess.check_output(
            ['nvidia-smi', f'--query-gpu={field}',
                '--format=csv,nounits,noheader'], 
            encoding='utf-8'))
    def to_int(result):
        return int(result.strip().split('\n')[0])
    
    used = to_int(query('memory.used'))
    total = to_int(query('memory.total'))
    pct = used/total
    print('\n' + msg, f'{100*pct:2.1f}% ({used} out of {total})')    


# ### Running the Training Loop

def get_training_obj(params):
    config = RobertaConfig.from_pretrained(params['model_name'], num_labels=2)
    model = RobertaForSequenceClassification.from_pretrained(params['model_name'], config=config).to(params['device'])
    no_decay = ['bias', 'LayerNorm.weight']
    gpd_params = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 
         'weight_decay': params['weight_decay']},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 
         'weight_decay': 0.0}
        ]
    optimizer = AdamW(gpd_params, lr=params['lr'], eps=params['adam_epsilon'])
    warmup_steps = int(params['warmup_ratio'] * params['total_steps'])
    scheduler = WarmupLinearSchedule(optimizer, warmup_steps=warmup_steps, t_total=params['total_steps'])
    return model, optimizer, scheduler


show_gpu('Initial GPU memory usage:')
for i in range(2):
    model, optimizer, scheduler = get_training_obj(params)
    show_gpu(f'{i}: GPU memory usage after loading training objects:')
    for epoch in range(1):
        epoch_start = time.time()
        model.train()
        for batch in dp.train_dataloader:
            xb,mb,_,yb = tuple(t.to(params['device']) for t in batch)
            outputs = model(input_ids = xb, attention_mask = mb, labels = yb)
            loss = outputs[0]
            loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
    show_gpu(f'{i}: GPU memory usage after training model:')
    del model, optimizer, scheduler, loss, outputs
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    show_gpu(f'{i}: GPU memory usage after clearing cache:')