| 34 | |
| 35 | |
| 36 | class Engine(object): |
| 37 | def __init__(self, custom_parser=None): |
| 38 | logger.info("PyTorch Version {}".format(torch.__version__)) |
| 39 | self.state = State() |
| 40 | self.devices = None |
| 41 | self.distributed = False |
| 42 | |
| 43 | if custom_parser is None: |
| 44 | self.parser = argparse.ArgumentParser() |
| 45 | else: |
| 46 | assert isinstance(custom_parser, argparse.ArgumentParser) |
| 47 | self.parser = custom_parser |
| 48 | |
| 49 | self.inject_default_parser() |
| 50 | self.args = self.parser.parse_args() |
| 51 | |
| 52 | self.continue_state_object = self.args.continue_fpath |
| 53 | if "WORLD_SIZE" in os.environ: |
| 54 | self.distributed = int(os.environ["WORLD_SIZE"]) > 1 |
| 55 | print(self.distributed) |
| 56 | |
| 57 | if self.distributed: |
| 58 | # self.local_rank = self.args.local_rank |
| 59 | self.local_rank = int(os.environ["LOCAL_RANK"]) |
| 60 | self.world_size = int(os.environ["WORLD_SIZE"]) |
| 61 | torch.cuda.set_device(self.local_rank) |
| 62 | os.environ["MASTER_ADDR"] = "127.0.0.1" |
| 63 | # os.environ['MASTER_PORT'] = self.args.port |
| 64 | torch.distributed.init_process_group(backend="nccl") |
| 65 | print(self.local_rank) |
| 66 | self.devices = [0, 1] # [i for i in range(self.world_size)] |
| 67 | else: |
| 68 | self.local_rank = int(os.environ["LOCAL_RANK"]) |
| 69 | self.devices = [0, 1] # parse_devices(self.args.devices) |
| 70 | |
| 71 | self.checkpoint_state = [] |
| 72 | |
| 73 | def inject_default_parser(self): |
| 74 | p = self.parser |
| 75 | p.add_argument("-d", "--devices", default="", help="set data parallel training") |
| 76 | p.add_argument( |
| 77 | "-c", |
| 78 | "--continue", |
| 79 | type=extant_file, |
| 80 | metavar="FILE", |
| 81 | dest="continue_fpath", |
| 82 | help="continue from one certain checkpoint", |
| 83 | ) |
| 84 | p.add_argument("--local_rank", default=0, type=int, help="process rank on node") |
| 85 | p.add_argument( |
| 86 | "-p", |
| 87 | "--port", |
| 88 | type=str, |
| 89 | default="16005", |
| 90 | dest="port", |
| 91 | help="port for init_process_group", |
| 92 | ) |
| 93 | |