imnet_evaluate/train.py [79:160]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - @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. """ model: nn.Module def save(self, filename: str) -> None: data = attr.asdict(self) # store only the state dict data["model"] = self.model.state_dict() torch.save(data, filename) @classmethod def load(cls, filename: str, default: "TrainerState") -> "TrainerState": data = torch.load(filename) # We need this default to load the state dict model = default.model model.load_state_dict(data["model"]) data["model"] = model return cls(**data) class Trainer: def __init__(self, train_cfg: TrainerConfig, cluster_cfg: ClusterConfig) -> None: self._train_cfg = train_cfg self._cluster_cfg = cluster_cfg def __call__(self) -> Optional[float]: """ Called 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 checkpoint(self, rm_init=True): save_dir = osp.join(self._train_cfg.save_folder, str(self._train_cfg.job_id)) os.makedirs(save_dir, exist_ok=True) self._state.save(osp.join(save_dir, "checkpoint.pth")) self._state.save(osp.join(save_dir, "checkpoint_"+str(self._state.epoch)+".pth")) if rm_init: os.remove(self._cluster_cfg.dist_url[7:]) empty_trainer = Trainer(self._train_cfg, self._cluster_cfg) return empty_trainer def _setup_process_group(self) -> None: torch.cuda.set_device(self._train_cfg.local_rank) torch.distributed.init_process_group( backend=self._cluster_cfg.dist_backend, init_method=self._cluster_cfg.dist_url, world_size=self._train_cfg.num_tasks, rank=self._train_cfg.global_rank, ) print(f"Process group: {self._train_cfg.num_tasks} tasks, rank: {self._train_cfg.global_rank}") def _init_state(self) -> None: """ Initialize the state and load it from an existing checkpoint if any """ torch.manual_seed(0) np.random.seed(0) print("Create data loaders", flush=True) print("Input size : "+str(self._train_cfg.input_size)) print("Model : " + str(self._train_cfg.architecture) ) backbone_architecture=None if self._train_cfg.architecture=='PNASNet' : backbone_architecture='pnasnet5large' transformation=get_transforms(input_size=self._train_cfg.input_size,test_size=self._train_cfg.input_size, kind='full', crop=True, need=('train', 'val'), backbone=backbone_architecture) transform_test = transformation['val'] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - imnet_extract/train.py [27:112]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - @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. """ model: nn.Module def save(self, filename: str) -> None: data = attr.asdict(self) # store only the state dict data["model"] = self.model.state_dict() torch.save(data, filename) @classmethod def load(cls, filename: str, default: "TrainerState") -> "TrainerState": data = torch.load(filename) # We need this default to load the state dict model = default.model model.load_state_dict(data["model"]) data["model"] = model return cls(**data) class Trainer: def __init__(self, train_cfg: TrainerConfig, cluster_cfg: ClusterConfig) -> None: self._train_cfg = train_cfg self._cluster_cfg = cluster_cfg def __call__(self) -> Optional[float]: """ Called 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 checkpoint(self, rm_init=True): # will be called by submitit in case of preemption save_dir = osp.join(self._train_cfg.save_folder, str(self._train_cfg.job_id)) os.makedirs(save_dir, exist_ok=True) self._state.save(osp.join(save_dir, "checkpoint.pth")) self._state.save(osp.join(save_dir, "checkpoint_"+str(self._state.epoch)+".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: 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 empty_trainer def _setup_process_group(self) -> None: torch.cuda.set_device(self._train_cfg.local_rank) torch.distributed.init_process_group( backend=self._cluster_cfg.dist_backend, init_method=self._cluster_cfg.dist_url, world_size=self._train_cfg.num_tasks, rank=self._train_cfg.global_rank, ) print(f"Process group: {self._train_cfg.num_tasks} tasks, rank: {self._train_cfg.global_rank}") def _init_state(self) -> None: """ Initialize the state and load it from an existing checkpoint if any """ torch.manual_seed(0) np.random.seed(0) print("Create data loaders", flush=True) print("Input size : "+str(self._train_cfg.input_size)) print("Model : " + str(self._train_cfg.architecture) ) backbone_architecture=None if self._train_cfg.architecture=='PNASNet' : backbone_architecture='pnasnet5large' transformation=get_transforms(input_size=self._train_cfg.input_size,test_size=self._train_cfg.input_size, kind='full', crop=True, need=('train', 'val'), backbone=backbone_architecture) transform_test = transformation['val'] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -