main.py

from __future__ import absolute_import, division, print_function

import os
import sys
import glob
import time
import json
import torch
import pickle
import random
import logging
import subprocess
import numpy as np
import jsbeautifier
from args import load_args
from logger import create_logger
from train import baseline_training
from evaluate import baseline_eval, FewShotBenchmark
from utils import set_seed, dist_training, backup_codes_to_project_folder

opts = jsbeautifier.default_options()

def main():

    start_time = time.time()
    args = load_args()
    if args.load_args is not None:
        args = torch.load(args.load_args)

    logger_id = args.logger_id if args.logger_id is not None else max(args.local_rank, 0)
    logger = create_logger(
                os.path.join(args.output_dir, "{}-{}.log".format(args.task_name, logger_id)), 
                process_rank = logger_id
            )
    args.device = str(args.device) if  hasattr(args,'device') else None
    logger.info("{}".format(jsbeautifier.beautify(json.dumps(args.__dict__), opts)))

    args = dist_training(args)
    
    set_seed(args.seed)
    logger.warning(
        " Process rank: {}, Python Version : {}\n"
        " torch version {},  cuda version: {},  CuDNN version : {}\n"
        " device: {},  n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
            args.local_rank,
            sys.version.replace("\n", " "),
            torch.__version__,
            str(torch.version.cuda),
            torch.backends.cudnn.version(),
            args.device,
            args.n_gpu,
            bool(args.local_rank != -1),
            args.fp16,
        )
    )
    
    if args.do_eval:
        baseline_eval(args, logger)

    if args.do_few_shot_benchamrk:
        FewShotBenchmark(args, logger)

    if args.do_train:
        baseline_training(args, logger)

    end_time = time.time()
    ty_res = time.gmtime(end_time-start_time)
    logger.info("Total runtime : {}".format( time.strftime("%H:%M:%S",ty_res) ))

if __name__ == "__main__":
    
    main()