| | |
|---|
| | | |
|---|
| | | os.makedirs(os.path.join(self.output_dir, "tensorboard"), exist_ok=True) |
|---|
| | | self.writer = SummaryWriter(os.path.join(self.output_dir, "tensorboard")) if rank == 0 else None |
|---|
| | | |
|---|
| | | |
|---|
| | | |
|---|
| | | def _save_checkpoint(self, epoch, step=None): |
|---|
| | |
|---|
| | | """ |
|---|
| | | ckpt = os.path.join(resume_path, "model.pt") |
|---|
| | | if os.path.isfile(ckpt): |
|---|
| | | checkpoint = torch.load(ckpt) |
|---|
| | | checkpoint = torch.load(ckpt, map_location="cpu") |
|---|
| | | self.start_epoch = checkpoint['epoch'] + 1 |
|---|
| | | # self.model.load_state_dict(checkpoint['state_dict']) |
|---|
| | | src_state = checkpoint['state_dict'] |
|---|
| | |
|---|
| | | print(f"Checkpoint loaded successfully from '{ckpt}'") |
|---|
| | | else: |
|---|
| | | print(f"No checkpoint found at '{ckpt}', does not resume status!") |
|---|
| | | |
|---|
| | | |
|---|
| | | self.model.to(self.device) |
|---|
| | | if self.use_ddp or self.use_fsdp: |
|---|
| | | dist.barrier() |
|---|
| | | |
|---|
