File size: 2,617 Bytes
12fea4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
import logging
from pathlib import Path
import torch
from cas9.data import DataState
from torch import nn
from torch.optim import Optimizer
class TrainState:
def __init__(
self,
model: nn.Module,
optimizer: Optimizer,
step: int,
# data_state: DataState,
):
self._model = model
self._optimizer = optimizer
self._step = step
# self._data_state = data_state
@property
def step(self) -> int:
return self._step
@step.setter
def step(self, value: int) -> None:
self._step = value
@property
def optimizer(self) -> Optimizer:
return self._optimizer
@property
def model(self) -> nn.Module:
return self._model
# @property
# def data_state(self) -> DataState:
# return self._data_state
def compile_model(self) -> None:
self._model = torch.compile(self._model)
def restore_checkpoint(
self, ckpt_dir: Path, device: torch.device, rank: int
) -> None:
if ckpt_dir.exists():
loaded_state = torch.load(ckpt_dir, map_location=device, weights_only=True)
self.optimizer.load_state_dict(loaded_state["optimizer"])
self.model.module.load_state_dict(loaded_state["model"])
self.step = loaded_state["step"]
# self._data_state.test.load_state_dict(loaded_state["test_sampler"])
# self._data_state.train.sampler.load_state_dict(
# loaded_state["train_sampler"]
# )
else:
ckpt_dir.parent.mkdir(exist_ok=True, parents=True)
if rank == 0:
logging.warning(
f"No checkpoint found at {ckpt_dir}. Returned the same state as input"
)
def save_checkpoint(self, ckpt_dir: str, rank: int) -> None:
saved_state = {
"optimizer": self.optimizer.state_dict(),
"model": self.model.module.state_dict(),
"step": self.step,
# "train_sampler": self._data_state.train.sampler.state_dict(),
# "test_sampler": self._data_state.test.sampler.state_dict(),
}
if rank == 0:
torch.save(saved_state, ckpt_dir)
def eval(self) -> None:
self.train(training=False)
def train(self, training: bool = True) -> None:
self._model.train(mode=training)
|