def __init__(self): self.parser = argparse.ArgumentParser() self.parser.add_argument('--gin', default = 'poem/poem.gin') self.opt = self.parser.parse_args() gInit(self.opt) self.opt.manualSeed = ginM('manualSeed') self.opt.port = ginM('port') if OnceExecWorker: rSeed(self.opt.manualSeed) self.opt.num_gpu = torch.cuda.device_count() if pO.HVD: self.opt.world_size = hvd.size() self.opt.rank = hvd.rank()
def gInit(opt): global pO, OnceExecWorker gin.parse_config_file(opt.gin) pO = parOptions(**{ginM('dist'): True}) if pO.HVD: hvd.init() torch.cuda.set_device(hvd.local_rank()) OnceExecWorker = (pO.HVD and hvd.rank() == 0) or (pO.DP) cudnn.benchmark = True
opt.world_size = opt.num_gpu opt.rank = rank train(opt) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--gin', help='Gin config file') opt = parser.parse_args() gInit(opt) print('[2] FINISHED GIN INIT') opt.manualSeed = ginM('manualSeed') opt.port = ginM('port') if OnceExecWorker: rSeed(opt.manualSeed) opt.num_gpu = torch.cuda.device_count() if pO.HVD: opt.world_size = hvd.size() opt.rank = hvd.rank() print('[3] SETTINGS DONE; INIT TRAIN') if not pO.DDP: train(opt) else: