File size: 961 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import yaml
import argparse
from trainer import ExpMultiGpuTrainer
def arg_parser():
parser = argparse.ArgumentParser(description="config")
parser.add_argument("--config",
type=str,
default="config/Recce.yml",
help="Specified the path of configuration file to be used.")
parser.add_argument("--local_rank", default=0,
type=int,
help="Specified the node rank for distributed training.")
return parser.parse_args()
if __name__ == '__main__':
import torch
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True
arg = arg_parser()
config = arg.config
with open(config) as config_file:
config = yaml.load(config_file, Loader=yaml.FullLoader)
config["config"]["local_rank"] = arg.local_rank
trainer = ExpMultiGpuTrainer(config, stage="Train")
trainer.train()
|