| from pathlib import Path |
| import torch |
| import pickle |
| import argparse |
| import logging |
| import torch.distributed as dist |
| from config import MyParser |
| from steps import trainer |
|
|
|
|
| if __name__ == "__main__": |
| formatter = ( |
| "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s" |
| ) |
| logging.basicConfig(format=formatter, level=logging.INFO) |
| |
| torch.cuda.empty_cache() |
| args = MyParser().parse_args() |
| logging.info(args) |
| exp_dir = Path(args.exp_dir) |
| exp_dir.mkdir(exist_ok=True, parents=True) |
| logging.info(f"exp_dir: {str(exp_dir)}") |
|
|
| if args.resume: |
| resume = args.resume |
| assert(bool(args.exp_dir)) |
| with open("%s/args.pkl" % args.exp_dir, "rb") as f: |
| old_args = pickle.load(f) |
| new_args = vars(args) |
| old_args = vars(old_args) |
| for key in new_args: |
| if key not in old_args or old_args[key] != new_args[key]: |
| old_args[key] = new_args[key] |
| args = argparse.Namespace(**old_args) |
| args.resume = resume |
| else: |
| with open("%s/args.pkl" % args.exp_dir, "wb") as f: |
| pickle.dump(args, f) |
|
|
| dist.init_process_group(backend='nccl', init_method='env://') |
| rank = dist.get_rank() |
| world_size = dist.get_world_size() |
| torch.cuda.set_device(rank) |
| my_trainer = trainer.Trainer(args, world_size, rank) |
| my_trainer.train() |
