mdr/retrieval/mhop_trainer.py [38:146]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - apex.amp.register_half_function(torch, 'einsum') from apex import amp @attr.s(auto_attribs=True) class TrainerState: """ Contains the state of the Trainer. It can be saved to checkpoint the training and loaded to resume it. """ epoch: int model: nn.Module optimizer: optim.Optimizer lr_scheduler: torch.optim.lr_scheduler._LRScheduler global_step: int def save(self, filename: str) -> None: data = attr.asdict(self) # store only the state dict data["model"] = self.model.state_dict() data["optimizer"] = self.optimizer.state_dict() data["lr_scheduler"] = self.lr_scheduler.state_dict() torch.save(data, filename) @classmethod def load(cls, filename: str, default: "TrainerState", gpu: int) -> "TrainerState": data = torch.load(filename, map_location=lambda storage, loc: storage.cuda(gpu)) # We need this default to load the state dict model = default.model model.load_state_dict(data["model"]) data["model"] = model optimizer = default.optimizer optimizer.load_state_dict(data["optimizer"]) data["optimizer"] = optimizer lr_scheduler = default.lr_scheduler lr_scheduler.load_state_dict(data["lr_scheduler"]) data["lr_scheduler"] = lr_scheduler return cls(**data) class Trainer: def __init__(self, train_cfg: NamedTuple, cluster_cfg: ClusterConfig) -> None: self._train_cfg = train_cfg self._cluster_cfg = cluster_cfg def __call__(self) -> Optional[float]: """ Called by submitit for each task. :return: The master task return the final accuracy of the model. """ self._setup_process_group() self._init_state() final_acc = self._train() return final_acc def log(self, log_data: dict): job_env = submitit.JobEnvironment() # z = {**vars(self._train_cfg), **log_data} save_dir = Path(self._train_cfg.output_dir) os.makedirs(save_dir, exist_ok=True) with open(save_dir / 'log.txt', 'a') as f: f.write(json.dumps(log_data) + '\n') def checkpoint(self, rm_init=True) -> submitit.helpers.DelayedSubmission: # will be called by submitit in case of preemption job_env = submitit.JobEnvironment() save_dir = osp.join(self._train_cfg.output_dir, str(job_env.job_id)) os.makedirs(save_dir, exist_ok=True) self._state.save(osp.join(save_dir, "checkpoint.pth")) # Trick here: when the job will be requeue, we will use the same init file # but it must not exist when we initialize the process group # so we delete it, but only when this method is called by submitit for requeue if rm_init and osp.exists(self._cluster_cfg.dist_url[7:]): os.remove(self._cluster_cfg.dist_url[7:]) # remove file:// at the beginning # This allow to remove any non-pickable part of the Trainer instance. empty_trainer = Trainer(self._train_cfg, self._cluster_cfg) return submitit.helpers.DelayedSubmission(empty_trainer) def _setup_process_group(self) -> None: job_env = submitit.JobEnvironment() torch.cuda.set_device(job_env.local_rank) torch.distributed.init_process_group( backend=self._cluster_cfg.dist_backend, init_method=self._cluster_cfg.dist_url, world_size=job_env.num_tasks, rank=job_env.global_rank, ) print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}") def _init_state(self) -> None: """ Initialize the state and load it from an existing checkpoint if any """ job_env = submitit.JobEnvironment() if job_env.global_rank == 0: # config_path = Path(args.save_folder) / str(job_env.job_id) / 'config.json' os.makedirs(self._train_cfg.output_dir, exist_ok=True) config_path = Path(self._train_cfg.output_dir) / 'config.json' with open(config_path, "w") as g: g.write(json.dumps(self._train_cfg._asdict())) print(f"Setting random seed {self._train_cfg.seed}", flush=True) random.seed(self._train_cfg.seed) np.random.seed(self._train_cfg.seed) torch.manual_seed(self._train_cfg.seed) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - mdr/retrieval/single_trainer.py [37:145]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - apex.amp.register_half_function(torch, 'einsum') from apex import amp @attr.s(auto_attribs=True) class TrainerState: """ Contains the state of the Trainer. It can be saved to checkpoint the training and loaded to resume it. """ epoch: int model: nn.Module optimizer: optim.Optimizer lr_scheduler: torch.optim.lr_scheduler._LRScheduler global_step: int def save(self, filename: str) -> None: data = attr.asdict(self) # store only the state dict data["model"] = self.model.state_dict() data["optimizer"] = self.optimizer.state_dict() data["lr_scheduler"] = self.lr_scheduler.state_dict() torch.save(data, filename) @classmethod def load(cls, filename: str, default: "TrainerState", gpu: int) -> "TrainerState": data = torch.load(filename, map_location=lambda storage, loc: storage.cuda(gpu)) # We need this default to load the state dict model = default.model model.load_state_dict(data["model"]) data["model"] = model optimizer = default.optimizer optimizer.load_state_dict(data["optimizer"]) data["optimizer"] = optimizer lr_scheduler = default.lr_scheduler lr_scheduler.load_state_dict(data["lr_scheduler"]) data["lr_scheduler"] = lr_scheduler return cls(**data) class Trainer: def __init__(self, train_cfg: NamedTuple, cluster_cfg: ClusterConfig) -> None: self._train_cfg = train_cfg self._cluster_cfg = cluster_cfg def __call__(self) -> Optional[float]: """ Called by submitit for each task. :return: The master task return the final accuracy of the model. """ self._setup_process_group() self._init_state() final_acc = self._train() return final_acc def log(self, log_data: dict): job_env = submitit.JobEnvironment() # z = {**vars(self._train_cfg), **log_data} save_dir = Path(self._train_cfg.output_dir) os.makedirs(save_dir, exist_ok=True) with open(save_dir / 'log.txt', 'a') as f: f.write(json.dumps(log_data) + '\n') def checkpoint(self, rm_init=True) -> submitit.helpers.DelayedSubmission: # will be called by submitit in case of preemption job_env = submitit.JobEnvironment() save_dir = osp.join(self._train_cfg.output_dir, str(job_env.job_id)) os.makedirs(save_dir, exist_ok=True) self._state.save(osp.join(save_dir, "checkpoint.pth")) # Trick here: when the job will be requeue, we will use the same init file # but it must not exist when we initialize the process group # so we delete it, but only when this method is called by submitit for requeue if rm_init and osp.exists(self._cluster_cfg.dist_url[7:]): os.remove(self._cluster_cfg.dist_url[7:]) # remove file:// at the beginning # This allow to remove any non-pickable part of the Trainer instance. empty_trainer = Trainer(self._train_cfg, self._cluster_cfg) return submitit.helpers.DelayedSubmission(empty_trainer) def _setup_process_group(self) -> None: job_env = submitit.JobEnvironment() torch.cuda.set_device(job_env.local_rank) torch.distributed.init_process_group( backend=self._cluster_cfg.dist_backend, init_method=self._cluster_cfg.dist_url, world_size=job_env.num_tasks, rank=job_env.global_rank, ) print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}") def _init_state(self) -> None: """ Initialize the state and load it from an existing checkpoint if any """ job_env = submitit.JobEnvironment() if job_env.global_rank == 0: # config_path = Path(args.save_folder) / str(job_env.job_id) / 'config.json' os.makedirs(self._train_cfg.output_dir, exist_ok=True) config_path = Path(self._train_cfg.output_dir) / 'config.json' with open(config_path, "w") as g: g.write(json.dumps(self._train_cfg._asdict())) print(f"Setting random seed {self._train_cfg.seed}", flush=True) random.seed(self._train_cfg.seed) np.random.seed(self._train_cfg.seed) torch.manual_seed(self._train_cfg.seed) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -