train.py

import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import sys
import random
import time
from datetime import datetime
import numpy as np
import math
import argparse
# random.seed(42)
from tqdm import tqdm

import logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(message)s")

import torch

import models, configs, data_loader 
from modules import get_cosine_schedule_with_warmup
from utils import similarity, normalize
from data_loader import *

    
def train(args):
    fh = logging.FileHandler(f"./output/{args.model}/{args.dataset}/logs.txt")
                                      # create file handler which logs even debug messages
    logger.addHandler(fh)# add the handlers to the logger
    timestamp = datetime.now().strftime('%Y%m%d%H%M') 
    
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu") 
    
    def save_model(model, epoch):
        torch.save(model.state_dict(), f'./output/{args.model}/{args.dataset}/models/epo{epoch}.h5')

    def load_model(model, epoch, to_device):
        assert os.path.exists(f'./output/{args.model}/{args.dataset}/models/epo{epoch}.h5'), f'Weights at epoch {epoch} not found'
        model.load_state_dict(torch.load(f'./output/{args.model}/{args.dataset}/models/epo{epoch}.h5', map_location=to_device))

    config = getattr(configs, 'config_'+args.model)()
    print(config)
    
    ###############################################################################
    # Load data
    ###############################################################################
    data_path = args.data_path+args.dataset+'/'
    train_set = eval(config['dataset_name'])(config, data_path, 
                                  config['train_ir'], config['n_node'],
                                  config['train_desc'], config['desc_len'])
    
    data_loader = torch.utils.data.DataLoader(dataset=train_set, batch_size=config['batch_size'],
                                       shuffle=True, drop_last=False, num_workers=1, pin_memory=True)
    # os.environ['MASTER_ADDR'] = 'localhost'
    # os.environ['MASTER_PORT'] = '5679'
    # torch.cuda.set_device(args.local_rank)
    # torch.distributed.init_process_group(backend='nccl', init_method=args.init_method, world_size=torch.cuda.device_count(), rank = args.local_rank)
    # device = torch.device("cuda", args.local_rank)
    # from torch.utils.data.distributed import DistributedSampler
    # from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler, TensorDataset
    # train_sampler = DistributedSampler(train_set)
    # data_loader = DataLoader(train_set, sampler=train_sampler, batch_size=args.train_batch_size*4, num_workers=1,
    #                               pin_memory=False)

    ###############################################################################
    # Define the models
    ###############################################################################
    logger.info('Constructing Model..')
    model = getattr(models, args.model)(config) #initialize the model
    if args.reload_from>0:
        load_model(model, args.reload_from, device)    
    logger.info('done')
    model.to(device)
    # model = torch.nn.DataParallel(model)

    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]    
    # origin: AdamW
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=config['learning_rate'], eps=config['adam_epsilon'])        
    # no scheduler in paper's original code
    scheduler = get_cosine_schedule_with_warmup(
            optimizer, num_warmup_steps=config['warmup_steps'], 
            num_training_steps=len(data_loader)*config['nb_epoch']) # do not foget to modify the number when dataset is changed
    if config['fp16']:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=config['fp16_opt_level'])
        
    print('---model parameters---')
    num_params = 0
    for param in model.parameters():
        num_params += param.numel()
    print(num_params / 1e6)
    
    n_iters = len(data_loader)
    itr_global = args.reload_from+1 
    for epoch in range(int(args.reload_from)+1, config['nb_epoch']+1): 
        itr_start_time = time.time()
        losses=[]

        # delta0 = 0
        # delta1 = 0
        # delta2 = 0
        # delta3 = 0
        for batch in data_loader:
            # clock1 = time.time()
            
            model.train()
            batch_gpu = [tensor.to(device) for tensor in batch]
            loss = model(*batch_gpu)
            
            # clock2 = time.time()
            # delta0 += clock2 - clock1
            
            if config['fp16']:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
                torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 5.0)
            else:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            
            # clock3 = time.time()
            # delta1 += clock3 - clock2

            optimizer.step()
            scheduler.step()
            model.zero_grad()
            
            # clock4 = time.time()
            # delta2 += clock4 - clock3

            losses.append(loss.item())
            
            if itr_global % args.log_every == 0:
                elapsed = time.time() - itr_start_time
                logger.info('epo:[%d/%d] itr:[%d/%d] step_time:%ds Loss=%.5f'%
                        (epoch, config['nb_epoch'], itr_global%n_iters, n_iters, elapsed, np.mean(losses)))
                # logger.info('%ds %ds %ds'%(delta0, delta1, delta2) )
                losses=[] 
                itr_start_time = time.time() 
                # delta0 = 0
                # delta1 = 0
                # delta2 = 0
                # delta3 = 0
            itr_global = itr_global + 1
        
        # save every epoch
        if epoch >= 10:
            if epoch % 5 == 0:
                save_model(model, epoch)
        

def parse_args():
    parser = argparse.ArgumentParser("Train and Validate The Code Search (Embedding) Model")
    parser.add_argument('--data_path', type=str, default='./data/', help='location of the data corpus')
    parser.add_argument('--model', type=str, default='IREmbeder', help='model name')
    parser.add_argument('--dataset', type=str, default='c_python', help='name of dataset.java, python')
    parser.add_argument('--reload_from', type=int, default=-1, help='epoch to reload from')
   
    parser.add_argument('-g', '--gpu_id', type=int, default=0, help='GPU ID')
    parser.add_argument('-v', "--visual",action="store_true", default=False, help="Visualize training status in tensorboard")
    # Training Arguments
    parser.add_argument('--log_every', type=int, default=100, help='interval to log autoencoder training results')
    parser.add_argument('--seed', type=int, default=1111, help='random seed')  # 1111

    parser.add_argument('--init_method', default='env://', type=str)
    parser.add_argument('--local_rank', type=int, default=0)

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    
    # make output directory if it doesn't already exist
    os.makedirs(f'./output/{args.model}/{args.dataset}/models', exist_ok=True)
    os.makedirs(f'./output/{args.model}/{args.dataset}/tmp_results', exist_ok=True)
    
    torch.backends.cudnn.benchmark = True # speed up training by using cudnn
    torch.backends.cudnn.deterministic = True # fix the random seed in cudnn
   
    train(args)