File size: 2,585 Bytes
7f316fe | 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 | # 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 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)
|