| 19 | |
| 20 | |
| 21 | class TrainState(object): |
| 22 | def __init__(self, optimizer, step, model=None, model_ema=None): |
| 23 | self.optimizer = optimizer |
| 24 | self.step = step |
| 25 | self.model = model |
| 26 | self.model_ema = model_ema |
| 27 | |
| 28 | def ema_update(self, rate=0.9999): |
| 29 | if self.model_ema is not None: |
| 30 | ema(self.model_ema, self.model, rate) |
| 31 | |
| 32 | def save(self, path): |
| 33 | os.makedirs(path, exist_ok=True) |
| 34 | torch.save(self.step, os.path.join(path, "step.pth")) |
| 35 | for key, val in self.__dict__.items(): |
| 36 | if key != "step" and val is not None: |
| 37 | torch.save(val.state_dict(), os.path.join(path, f"{key}.pth")) |
| 38 | |
| 39 | def load(self, path): |
| 40 | logging.info(f"load from {path}") |
| 41 | self.step = torch.load(os.path.join(path, "step.pth")) |
| 42 | for key, val in self.__dict__.items(): |
| 43 | if key != "step" and val is not None: |
| 44 | val.load_state_dict( |
| 45 | torch.load(os.path.join(path, f"{key}.pth"), map_location="cpu") |
| 46 | ) |
| 47 | |
| 48 | def resume(self, ckpt_root, step=None): |
| 49 | if not os.path.exists(ckpt_root): |
| 50 | return |
| 51 | if step is None: |
| 52 | ckpts = list(filter(lambda x: ".ckpt" in x, os.listdir(ckpt_root))) |
| 53 | if not ckpts: |
| 54 | return |
| 55 | steps = map(lambda x: int(x.split(".")[0]), ckpts) |
| 56 | step = max(steps) |
| 57 | ckpt_path = os.path.join(ckpt_root, f"{step}.ckpt") |
| 58 | logging.info(f"resume from {ckpt_path}") |
| 59 | self.load(ckpt_path) |
| 60 | |
| 61 | def to(self, device): |
| 62 | for key, val in self.__dict__.items(): |
| 63 | if isinstance(val, nn.Module): |
| 64 | val.to(device) |
| 65 | |
| 66 | |
| 67 | def cnt_params(model): |
no outgoing calls
no test coverage detected