(self, custom_parser=None)
| 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 |
nothing calls this directly
no test coverage detected